Intermediate 8 min · March 09, 2026

PyTorch Neural Network — The forward() Layer Bug

Loss drops, validation accuracy stuck at 10% random? Layers in forward() create new weights each batch—optimizer updates old ones.

Naren · Founder

Plain-English first. Then code. Then the interview question.

About

⚡Quick Answer

nn.Module is the base class for all PyTorch models — define layers in __init__, data flow in forward
super().__init__() is mandatory — without it, layers and parameters are not registered and model.parameters() returns empty
model.to(device) moves all parameters to GPU in one atomic call — never manually move individual weights
Defining layers inside forward() creates new untrained weights every pass — the optimizer updates weights that are immediately discarded
state_dict saves only learnable parameters — smaller, portable, and version-independent compared to saving the full model
model.eval() disables Dropout and freezes BatchNorm running statistics — always call it before inference or validation

Building a neural network in PyTorch revolves around one central idea: subclassing nn.Module. You define layers in __init__ and the data flow in forward. PyTorch automatically tracks all parameters, moves them to GPU with a single .to(device) call, and integrates cleanly with torch.optim for gradient-based training.

The nn.Module design solves parameter management at scale. Without it, you would manually track thousands of weight matrices, move each to GPU individually, and implement gradient updates by hand. The module system handles all of this through a unified interface: model.parameters() returns every learnable tensor, model.state_dict() serializes the full learnable state, and model.to(device) moves everything atomically — no risk of a weight matrix left behind on CPU while the rest of the model runs on GPU.

The production failure pattern I see most consistently: developers define layers inside forward() instead of __init__. This creates new uninitialized weights on every forward pass. The optimizer updates weights from the previous pass that no longer exist — they were replaced by fresh random tensors when forward() ran again. Training loss can decrease slightly due to random variation, which masks the bug entirely. Validation accuracy stays at random chance. No error is raised. The model trains for 100 epochs and learns nothing.

What Is Building a Neural Network in PyTorch and Why Does It Exist?

Building a neural network in PyTorch is the process of defining a model by subclassing nn.Module — PyTorch's foundational abstraction for everything that involves learnable parameters. It was designed to solve a specific problem: managing the lifecycle of thousands to billions of weight tensors without building that infrastructure yourself every time you train a model.

The architectural separation at the core of nn.Module is deliberate and meaningful. __init__ defines the static structure — which layers exist, their input and output sizes, how they are named. forward defines the dynamic behavior — how a tensor flows through those layers during each call. This separation is what makes the rest of the system work: PyTorch can inspect the model structure without running data through it, serialize only the parameters independently of the forward logic, and move the entire model to GPU atomically with model.to(device).

The key mechanism underneath all of this is Python's __setattr__ override in nn.Module. When you write self.fc1 = nn.Linear(784, 128) in __init__, PyTorch intercepts that assignment, detects that nn.Linear is itself an nn.Module, and registers it in an internal _modules dictionary. When you write self.weight = nn.Parameter(torch.randn(10, 5)), PyTorch detects nn.Parameter and registers it in _parameters. These dictionaries are what model.parameters(), model.state_dict(), and model.to(device) iterate over. None of this works if you skip super().__init__() — the dictionaries are never created, the __setattr__ override is never installed, and every layer you assign to self is just a plain Python attribute that PyTorch cannot see.

The practical consequence at production scale: a model with 100M parameters that is partially on GPU and partially on CPU produces wrong outputs without raising errors. Parameter groups that the optimizer cannot reach do not update. model.parameters() returning fewer tensors than expected is always a registration bug — not a configuration issue.

For 2026 deployments, the nn.Module contract also integrates with torch.compile() — PyTorch's graph compilation path introduced in 2.0 and stabilized through 2.2 and beyond. A properly structured nn.Module compiles cleanly with torch.compile(model), producing kernel fusion and operator overlap that can reduce training time by 30-50% on modern A100 and H100 hardware without changing a line of model code. Models with operations that break the graph — .numpy() calls inside forward, Python data structures used conditionally — either fail to compile or fall back to eager mode silently.

Enterprise Persistence: Saving and Loading Forge Models

In a production environment, training a model is only part of the story. You need to persist it, version it, load it reliably six months later, and reproduce its inference behavior exactly. Getting this wrong has a specific failure mode that is not immediately obvious: you load a model, it runs inference without any errors, and it produces predictions — predictions that are quietly wrong because Dropout is still active or because you loaded weights into the wrong architecture without noticing.

The core persistence decision in PyTorch is between saving the full model object and saving only the state_dict. torch.save(model, path) uses Python's pickle to serialize the entire model — code, architecture, and weights together. torch.save(model.state_dict(), path) serializes only the learnable parameter tensors as an OrderedDict of name-to-tensor mappings. The state_dict approach is the production standard for three concrete reasons: the file is smaller because no Python code is embedded, it is portable because you can load weights into a model defined anywhere as long as the parameter names match, and it is safer because pickle can execute arbitrary code when deserializing, which is a real attack surface in shared model repositories.

The full checkpoint pattern extends this for training resumption. Saving only model.state_dict() is sufficient for inference deployment, but if you need to resume training from a checkpoint, you also need the optimizer state — Adam's moment estimates are not recomputed from scratch, and resuming without them produces different training dynamics than if training had never stopped. A complete checkpoint includes model state, optimizer state, epoch number, and the best validation metric so you know whether to update your best-model checkpoint.

One detail that bites teams in production: torch.load() defaults to weights_only=False in PyTorch versions before 2.4, which means it will execute arbitrary pickle code. In PyTorch 2.4+, the default changed to weights_only=True for state_dict loading, which is safer. If you are loading state_dicts — which you should be — explicitly pass weights_only=True regardless of version to future-proof your code and prevent security warnings in CI.

Containerizing the Forge Model Service

Getting a PyTorch model to run correctly on a developer workstation is step one. Getting it to run correctly in production — on a different machine, a different OS, a different GPU driver, possibly six months from now — is the actual engineering problem. Containerization with Docker is the standard answer, but the details matter more than most tutorials acknowledge.

The version pinning problem is where most teams make their first mistake. Pulling pytorch/pytorch:latest in production means your deployment environment changes every time a new PyTorch release ships. Changes between minor versions can affect numerical precision, change default behaviors for certain operations, and silently alter model outputs. Pin the full triple: PyTorch version, CUDA version, and cuDNN version. These three together determine the exact kernel implementations your model runs on. A mismatch between cuDNN versions on the same PyTorch base can produce numerically different outputs from the same weights.

The image size problem compounds quickly in multi-service deployments. A CUDA-enabled PyTorch runtime image is typically 5-7GB. A CPU-only image is under 1GB. If your inference service runs on CPU-optimized instances — which is common for cost efficiency in steady-state serving — you are pulling 5-7GB per node during deployments when 1GB would be sufficient. This is not a philosophical problem — it translates directly to longer deployment times, higher container registry egress costs, and slower autoscaling response.

The model weight inclusion problem is the third one. Baking a 500MB model file into a Docker image with COPY means every CI build, every image push, and every container pull moves that 500MB. For a team with 10 engineers committing multiple times a day, this accumulates. The correct pattern is to exclude model weights from the image and mount them from a volume, or download them at container startup from an object store like S3 or GCS. This keeps the image lean, makes weight updates independent of image rebuilds, and allows you to run canary deployments with different weight versions without rebuilding images.

Common Mistakes and How to Avoid Them

Most nn.Module bugs fall into a small set of categories. They are not obscure — they appear consistently across codebases from beginners and experienced engineers alike, usually under deadline pressure when someone is focused on getting the model working and skips a step that seemed optional.

Forgetting super().__init__() is the most foundational mistake, and it has a particularly frustrating failure mode: the error often does not surface immediately. You define your model, assign layers to self, and nothing explodes. The failure comes later when model.parameters() returns an empty iterator, model.to(device) does nothing, or torch.save(model.state_dict()) produces a file with zero keys. By that point, the developer is often deep into debugging the training loop rather than looking at model initialization.

Using Python lists to store layers is the mistake that catches experienced developers. If you have used other frameworks or written Python professionally, using a list of layers feels completely natural — it is idiomatic Python. But a Python list of nn.Module instances is invisible to PyTorch. The parameters in those layers are not in model.parameters(), they are not moved by model.to(device), and the optimizer cannot update them. The model runs, the loss changes slightly due to the layers in the list processing data, and nothing indicates the optimizer is completely ignoring them. Use nn.ModuleList for any list of modules, and nn.ModuleDict for any dictionary of named modules.

The .numpy() inside forward() mistake is common in teams transitioning from NumPy-heavy workflows. It always produces a RuntimeError if the tensor requires gradients, or a silent gradient chain break if you call .detach() first. Both are wrong inside forward(). All computation in forward() must stay in PyTorch tensor operations. If you need NumPy for debugging, do it outside the computation graph after calling .detach().cpu().

One 2026-specific addition worth calling out: with torch.compile() becoming the standard path for production training, any Python-level control flow in forward() that depends on tensor values — not tensor shapes, but actual data values — will prevent the compiler from tracing the graph cleanly. This was always a theoretical concern; now it is a practical one because compile() is in the default training stack for many teams. Keep forward() deterministic in its control flow — conditional branches should depend on constructor arguments, not on runtime tensor contents.

Manual Matrix Math vs PyTorch nn.Module

Aspect	Manual Matrix Math	PyTorch nn.Module
Parameter Tracking	Manual — you maintain a dict or list of weight tensors and must not forget any	Automatic — `model.parameters()` and `model.named_parameters()` iterate every registered tensor
GPU Portability	Manual — every tensor must be moved individually with .to(device), easy to miss one	Atomic — model.to(device) moves every registered parameter and buffer in a single call
Gradient Computation	Manual — you must call .backward() on the right tensor and implement update logic	Automatic — Autograd tracks the computation graph; torch.optim handles parameter updates
Model Serialization	Custom logic — you must know which tensors to save, in which order, and how to restore them	Built-in — `model.state_dict()` and `load_state_dict()` handle serialization with named keys
Training / Eval Mode	Manual — you must track mode state and toggle Dropout and BatchNorm behavior yourself	Built-in — `model.train()` and `model.eval()` propagate recursively to all child modules
Compiler Compatibility	None — manual tensor code has no structural guarantees for `torch.compile()` optimization	Full — properly structured nn.Module compiles cleanly with `torch.compile()` for 30-50% training speedup

Key Takeaways

Building a neural network in PyTorch means subclassing nn.Module — understanding what that abstraction provides (automatic parameter tracking, GPU portability, optimizer integration, serialization) is more important than memorizing the syntax.
super().__init__() is mandatory and must be the first line of every __init__ — without it, no layers are registered, model.parameters() returns empty, and model.to(device) does nothing.
Define layers in __init__, data flow in forward — this separation is the entire contract of nn.Module and violating it produces bugs that are silent, expensive to debug, and easy to prevent.
Use nn.ModuleList for lists of modules, nn.ModuleDict for named collections — Python lists and dicts are invisible to PyTorch's parameter tracking, serialization, and device management.
Call model(x), not model.forward(x) — the __call__ mechanism manages hooks, autograd tracking, and training/eval mode state that forward() alone does not.
model.eval() after loading weights is mandatory for inference — Dropout and BatchNorm behave fundamentally differently in training and eval mode, and the difference directly affects prediction quality.
Verify trainable parameter count immediately after model construction — any unexpected number indicates a registration bug that will waste training compute if left undetected.

Common Mistakes to Avoid

Forgetting to call super().__init__() in nn.Module subclass
Symptom: model.parameters() returns an empty iterator. Layers assigned to self are not registered. model.to(device) moves nothing. model.state_dict() returns an empty OrderedDict. Training runs without errors but no weights are being updated. The failure is silent until you check parameter count.
Fix: Add super().__init__() as the absolute first line of every __init__ method in every nn.Module subclass. This initializes the _parameters, _modules, _buffers, and _hooks internal dictionaries and installs PyTorch's __setattr__ override, which is what makes layer registration automatic. Without it, self.fc = nn.Linear(...) is just a plain Python attribute assignment.
Defining layers inside forward() instead of __init__()
Symptom: Training loss decreases slightly across epochs, giving the appearance of learning. Validation accuracy stays at random chance — 10% for 10 classes, 50% for binary classification. model.named_parameters() shows tensors, but their values change by only a negligible amount across training epochs. No errors are raised anywhere in the training loop.
Fix: Move every nn.Linear, nn.Conv2d, nn.BatchNorm, nn.Embedding, and any other learnable layer definition from forward() to __init__(). forward() should contain only tensor operations — calls to self.layer_name(x), activation functions, reshapes, and concatenations. Nothing that allocates parameters belongs there.
Using a Python list instead of nn.ModuleList for dynamic layers
Symptom: model.parameters() returns fewer parameters than expected — specifically, zero from the layers stored in the list. Layers in the list are not moved to GPU when model.to(device) is called, causing a device mismatch error on the first forward pass. model.state_dict() does not include those layers, so saving and loading the model silently drops them.
Fix: Replace every Python list of nn.Module instances with nn.ModuleList. Replace every Python dict of nn.Module instances with nn.ModuleDict. Both register their contents with PyTorch, making parameters visible to model.parameters(), moveable by model.to(device), and serializable by model.state_dict(). Verify the fix by checking parameter count immediately after model construction.
Calling model.forward(x) directly instead of model(x)
Symptom: Forward hooks registered with model.register_forward_hook() do not fire. Backward hooks registered with model.register_backward_hook() do not fire. Debugging and profiling tools that rely on hooks produce no output. In some configurations, autograd tracking setup is incomplete, causing subtle gradient computation issues.
Fix: Always call the model as a callable: output = model(input_tensor). Never call model.forward(input_tensor) directly. The nn.Module __call__ method triggers pre-forward hooks, manages autograd setup, invokes forward(), triggers post-forward hooks, and handles training/eval mode bookkeeping. Calling forward() directly bypasses all of this.
Converting tensors to NumPy inside forward()
Symptom: RuntimeError: Can't call numpy() on a Tensor that requires grad if the tensor is in the computation graph. Or a silent gradient chain break if .detach() is called first — the model runs, computes a loss, calls backward(), but the gradients are zero or None for all layers that produced the detached tensor.
Fix: Keep every computation inside forward() as a PyTorch tensor operation. PyTorch has equivalents for nearly every NumPy function — use them. If you genuinely need NumPy values for debugging or post-processing, do the conversion outside the model: output = model(x).detach().cpu().numpy(). The .detach() call must happen outside forward(), after the backward pass is complete.

Interview Questions on This Topic

QExplain why super().__init__() is non-negotiable in PyTorch. What happens internally to the _parameters and _modules dictionaries?Mid-levelReveal
nn.Module.__init__() initializes several internal dictionaries that are the foundation of the entire parameter management system: _parameters stores nn.Parameter objects (learnable weights and biases), _modules stores child nn.Module instances (sub-layers and sub-networks), _buffers stores non-parameter tensors like BatchNorm's running_mean and running_var, and _hooks stores registered forward and backward hook callbacks. When you write self.fc1 = nn.Linear(10, 5) in __init__, Python calls nn.Module's overridden __setattr__ method. This override inspects the assigned value — if it is an nn.Parameter, it goes into _parameters; if it is an nn.Module, it goes into _modules; otherwise it is a plain Python attribute. Without super().__init__(), these dictionaries are never created. The __setattr__ override is never installed. Every layer you assign to self becomes a plain Python attribute. model.parameters() returns an empty iterator, model.to(device) moves nothing, and model.state_dict() produces an empty dict — all silently, without errors.
QContrast nn.Module subclassing with nn.Sequential. In what specific architectural scenario is nn.Sequential technically impossible to use?SeniorReveal
nn.Sequential is a container that enforces a strictly linear data flow — the output of module N is passed as the sole input to module N+1. It cannot express: (1) residual or skip connections where the block input is added to or concatenated with the block output — y = F(x) + x requires access to x after it has passed through F, which Sequential cannot provide; (2) multiple inputs where two tensors enter the same module; (3) multiple outputs where one module produces two tensors for different downstream modules; (4) conditional branching in forward where the execution path changes based on input properties; (5) non-module operations between layers like tensor reshaping, dimension permutation, or selective masking. The canonical example is any ResNet block: output = conv(x) + x requires holding the original x while also computing conv(x), which Sequential's linear chaining cannot express. You must subclass nn.Module and write custom forward() for any architecture with these characteristics.
QDescribe the vanishing gradient problem. How does the choice of activation function in forward() mitigate it, and what is the relationship between this problem and residual connections?SeniorReveal
Vanishing gradients occur because backpropagation multiplies gradients through every layer via the chain rule. If each layer's local gradient is consistently less than 1.0 — which happens with saturating activations — the product shrinks exponentially with depth. After 20 layers with a per-layer gradient of 0.5, the gradient at the first layer is roughly 0.5^20, which is effectively zero. Early layers stop learning. Sigmoid saturates at both ends with gradient bounded by 0.25. Tanh saturates with gradient bounded by 1.0, but still shrinks. ReLU addresses this directly: its gradient is exactly 1.0 for positive inputs, so gradients pass through unchanged for active neurons. The failure mode of ReLU is dead neurons — units where the pre-activation is always negative get a gradient of zero and never recover. LeakyReLU uses a small gradient (0.01) for negative inputs to prevent permanent death. GELU, the smooth approximation used in transformers and modern MLPs, combines the non-saturation of ReLU with smooth gradients that train more stably. Residual connections solve the problem from a structural angle rather than an activation angle: y = F(x) + x means the gradient of y with respect to x is dF/dx + 1. The constant 1 ensures a gradient of at least 1.0 flows through the skip path regardless of how small dF/dx becomes — deep networks with residual connections do not vanish.
QWhat is the difference between model.parameters() and model.state_dict()? When would you use one over the other?Mid-levelReveal
model.parameters() returns an iterator over every nn.Parameter object registered in the module and all its children — learnable weights and biases only. It is the right input for optimizers: optimizer = torch.optim.Adam(model.parameters(), lr=1e-3). The optimizer needs the actual tensor objects so it can update them in place and maintain per-parameter state like Adam's first and second moment estimates. model.state_dict() returns an OrderedDict mapping string parameter names to their tensor values. It includes both learnable parameters and registered buffers — BatchNorm's running_mean, running_var, and num_batches_tracked are buffers, not parameters, but they are part of the model's inference state and must be saved and restored for correct behavior. This is the right interface for serialization: you save the dict, load it back, and the model produces identical outputs. Use parameters() for optimization — the optimizer needs live tensor references. Use state_dict() for persistence and deployment — it captures the complete inference state including non-learnable buffers that parameters() omits.
QHow does PyTorch TorchScript interact with a standard nn.Module, and what constraints does it impose on the forward() method for production deployment?SeniorReveal
TorchScript converts a dynamic nn.Module into a static, serializable graph that runs without a Python interpreter — enabling deployment to C++ runtimes, mobile devices, and latency-critical services. Two compilation paths exist: torch.jit.trace records the exact operations executed on a sample input. It is fast and produces an efficient graph, but it cannot capture Python control flow — any if/else or loop that branches differently for different inputs will be frozen to the branch taken during tracing. torch.jit.script analyzes Python source code and compiles it, handling control flow correctly, but requires type annotations on all inputs and outputs of forward() and restricts forward() to a subset of Python — no arbitrary Python objects, no dynamic attribute creation, no calls to non-TorchScript-compatible functions. The result in either case is a ScriptModule saved with torch.jit.save() and loaded anywhere LibTorch runs. For the nn.Module to compile cleanly: forward() must have static type annotations, all control flow must depend on types or constructor arguments rather than runtime tensor values, and all operations must be expressible in TorchScript's type system. With torch.compile() — the 2.0+ compilation path using Dynamo and Inductor — the constraints are looser: it handles more Python dynamism via graph breaks, falling back to eager for unsupported sections, while still compiling the static parts for kernel fusion and operator overlap gains.

Frequently Asked Questions

What is building a neural network in PyTorch in simple terms?

It is the process of defining a model's structure and behavior using PyTorch's nn.Module class. You write __init__ to declare which layers exist and how large they are. You write forward to describe how data moves through those layers to produce a prediction. PyTorch handles everything else: tracking the weights, computing gradients, moving parameters to GPU, and saving the trained model. You focus on the architecture. The framework handles the infrastructure.

Can I use multiple GPUs for my model?

Yes. PyTorch provides two approaches. nn.DataParallel wraps your model and splits each batch across multiple GPUs on a single machine — simpler to set up but has a known bottleneck at the parameter server on GPU 0 and does not scale well beyond 4 GPUs. DistributedDataParallel (DDP) runs a separate process per GPU, each with its own model replica, and synchronizes gradients via all-reduce after each backward pass — more setup required but scales linearly and is the production standard for multi-GPU training. For 2026 deployments, DDP with torch.compile() and mixed precision is the recommended training stack for serious model training on multi-GPU infrastructure.

What is the difference between a layer and a module in PyTorch?

Every layer in PyTorch — nn.Linear, nn.Conv2d, nn.BatchNorm1d, nn.Dropout — is itself a subclass of nn.Module. A module is the more general concept: it can be a single layer with a few parameters, or it can be a complex sub-network containing dozens of layers and other modules nested arbitrarily deep. When you build a model by subclassing nn.Module and assigning layers to self in __init__, your model is a module that contains other modules. The terms are used interchangeably in practice, but module is technically the correct term for any nn.Module subclass, while layer usually refers to a specific operation like a linear transformation or convolution.

Why do we use the forward method instead of just defining a __call__ method?

You define forward() because nn.Module's __call__ method calls forward() internally, but also wraps it with additional behavior that PyTorch needs: registering the forward pass with autograd for gradient tracking, firing any registered forward hooks (used by profilers, debuggers, and feature extraction tools), and managing training versus eval mode for layers like Dropout and BatchNorm. If you overrode __call__ directly, you would lose all of that. By defining forward() and calling the model as model(x), you get all the PyTorch infrastructure for free. This is why calling model.forward(x) directly — bypassing __call__ — is wrong even though it produces numerically identical output.

When should I use nn.ModuleList versus a Python list?

Use nn.ModuleList any time you have a collection of nn.Module instances that you want PyTorch to know about — which is essentially always. A Python list of layers is a plain Python object from PyTorch's perspective: the parameters inside those layers are not tracked by model.parameters(), not moved by model.to(device), not included in model.state_dict(), and not accessible to the optimizer. The model will run — Python will find the layers through the list — but the optimizer cannot update them and the weights are not saved when you checkpoint. Use nn.ModuleList for ordered collections of modules and nn.ModuleDict for named collections. If you only need to store hyperparameters or non-module configuration, a plain Python list or dict is fine.

🔥

That's PyTorch. Mark it forged?

8 min read · try the examples if you haven't