Introduction

nn.Module is the base class that every neural network model and layer in PyTorch is built from — providing the standard structure for organizing learnable parameters, defining a forward pass, and integrating seamlessly with Autograd, optimizers, and the rest of the PyTorch ecosystem. Whether you're using a single built-in layer like nn.Linear or building an entire custom architecture, everything ultimately inherits from and follows the same nn.Module pattern.

This is the direct PyTorch counterpart to the Keras layer/model system covered in the TensorFlow section — while the specific syntax differs, nn.Module serves fundamentally the same purpose: providing a consistent, reusable structure for defining neural network components.

Why Does nn.Module Matter?

nn.Module helps to:

• Provide a standard, consistent structure for defining any neural network component
• Automatically track all learnable parameters (weights and biases) within a model
• Integrate seamlessly with Autograd for gradient computation during training
• Support nesting — modules can contain other modules, building complex architectures from simple building blocks
• Enable easy saving, loading, and moving of models between CPU and GPU
• Serve as the foundation beneath every built-in PyTorch layer and every custom model

The Structure of an nn.Module

Whiteboard

Every custom model in PyTorch follows the same two-part pattern:

__init__: create and store the layers the model will use
          (this is also where you MUST call super().__init__())

forward: define the actual computation — how input data flows
         through the layers defined in __init__ to produce output

This mirrors the same __init__/call structure covered in the
Model Subclassing topic from the TensorFlow section, just with
PyTorch's own naming convention (forward instead of call).

A Basic Custom Model

Notice that layers (nn.Linear, nn.ReLU, nn.Sigmoid) are created once in __init__, while forward() defines the actual sequence of operations connecting them — you never call forward() directly; instead, calling model(input) automatically invokes it behind the scenes.

Common Built-In Layers

Automatic Parameter Tracking

Because every layer assigned as an attribute in __init__ is
itself an nn.Module, PyTorch automatically discovers and tracks
ALL learnable parameters across the entire model — this is
what allows a single call like optimizer = Adam(model.parameters())
to correctly gather every trainable weight and bias in the
model, regardless of how deeply nested the architecture is.

nn.Sequential: A Simpler Alternative for Simple Stacks

For simple, linear stacks of layers, nn.Sequential provides a more concise alternative to writing a full custom class — directly analogous to Keras's Sequential API covered in the TensorFlow section, offering the same convenience for straightforward architectures.

Nesting Modules Within Modules

This demonstrates a key strength of nn.Module: since a module can contain other modules as attributes, complex architectures can be built by composing smaller, reusable custom blocks — parameter tracking, device movement, and saving/loading all work correctly automatically, no matter how deeply nested the structure becomes.

Moving a Model Between CPU and GPU

Calling .to(device) on an nn.Module moves every parameter within it (including all nested sub-modules) to the specified device — the exact same .to() pattern covered in the earlier Tensors topic, applied here at the whole-model level.

Training Mode vs Evaluation Mode

nn.Module vs TensorFlow's Keras Layer/Model System

nn.Module Class vs nn.Sequential

Key Properties of nn.Module

• Every PyTorch layer and model inherits from nn.Module, providing a consistent structure and behavior.
• Layers are defined in __init__, and the forward pass logic is defined in forward().
• nn.Module automatically discovers and tracks all learnable parameters, including those in nested sub-modules.
• nn.Sequential offers a concise shortcut for simple, linear layer stacks, similar to Keras's Sequential API.
• .train() and .eval() explicitly control behavior for layers like Dropout and BatchNorm that differ between training and inference.

Where Is nn.Module Used?

Advantages

• Provides a single, consistent structure for both simple layers and entire complex models
• Automatic parameter tracking eliminates manual bookkeeping of weights across a model
• Supports arbitrary nesting, enabling clean, modular, reusable architecture design
• Seamlessly integrates with Autograd, optimizers, and device placement
• nn.Sequential offers convenience for simple cases without sacrificing full flexibility when needed

Limitations

• Requires understanding the __init__/forward() pattern and remembering to call super().__init__()
• Explicit .train()/.eval() mode switching is easy to forget, leading to subtle bugs
• More verbose than nn.Sequential for genuinely simple architectures
• Debugging deeply nested custom modules can require careful tracing through multiple forward() calls
• No built-in automatic shape inference the way some other frameworks provide (input sizes are typically specified explicitly)

Real-World Examples

Best Practices

• Always call super().__init__() as the first line inside a custom nn.Module's __init__ method.
• Use nn.Sequential for simple, linear architectures, and custom classes when more flexibility is needed.
• Remember to call model.train() before training and model.eval() before evaluation or inference.
• Build complex architectures by composing smaller, reusable custom nn.Module blocks.
• Use model.parameters() when setting up an optimizer to ensure all trainable weights are correctly included.

Interview Tip

A common interview question is:

"What is nn.Module, and why is calling super().__init__() necessary when defining a custom model?"

A strong answer is:

nn.Module is the base class that every PyTorch layer and model inherits from, providing automatic tracking of learnable parameters, integration with Autograd, and support for moving models between devices, saving/loading, and nesting sub-modules. Calling super().__init__() is necessary because it runs nn.Module's own internal setup code, which initializes the internal data structures — like the registry that tracks parameters and sub-modules — that PyTorch relies on to make features like model.parameters() and .to(device) work correctly; skipping this call would break that essential internal bookkeeping.

Explaining specifically what super().__init__() sets up makes your answer stronger.

Conclusion

nn.Module provides the foundational structure underlying every PyTorch model, standardizing how layers are defined, how parameters are tracked, and how the forward pass is computed, while integrating seamlessly with Autograd and the rest of the training pipeline. With nn.Module now covered, the next topic explores Optimizers — the algorithms responsible for actually using the gradients Autograd computes to update a model's parameters during training.