Training Trains Internship
Student login
← PyTorch learning path
PyTorch · Module 4 of 10

Building models with nn.Module

Define networks as Python classes.

📖 Key ideas

  1. Subclass nn.Module; define layers in __init__ and the computation in forward.
  2. Layers: nn.Linear, nn.Conv2d, nn.ReLU, nn.Dropout, nn.LSTM.
  3. nn.Sequential chains layers quickly.
  4. model.parameters() lists trainable weights.
  5. model.train() and model.eval() switch modes.
💡 Real-life example: nn.Module is a blueprint: __init__ lists the parts, forward describes how they connect.
nn.ModuleBase class for models
forward()Defines computation
nn.LinearFully connected layer
model.eval()Sets evaluation mode