How to create a custom nn.Module in PyTorch

Model code becomes easier to train, move between devices, and save when its layers belong to one torch.nn.Module object. PyTorch can then discover the component's parameters while the class keeps its tensor transformation behind a normal callable interface.

A child layer becomes registered when the module assigns it to an attribute in __init__(). The parent __init__() must run before those assignments, and forward() defines how input tensors pass through the registered children.

A small CPU regression model is enough to test both halves of that contract. Its layer names must appear in the module tree, while a real backward pass must produce gradients that let the optimizer change a weight.

1. Create custom_module_demo.py with the imports and registered child layers.

   custom_module_demo.py

   import torch
   from torch import nn
    
    
   torch.manual_seed(7)
    
    
   class SensorRegressor(nn.Module):
       def __init__(self, in_features=4, hidden_features=8, out_features=2):
           super().__init__()
           self.input = nn.Linear(in_features, hidden_features)
           self.activation = nn.ReLU()
           self.output = nn.Linear(hidden_features, out_features)

   super().init() initializes the module before input, activation, and output are registered as children.

2. Insert the forward method below the initializer inside SensorRegressor.

       def forward(self, features):
           hidden = self.activation(self.input(features))
           return self.output(hidden)

   Calling model(features) keeps the normal module dispatch path active, whereas a direct model.forward(features) call bypasses it.

3. Append the sample batch and training objects after the class definition.

   features = torch.tensor(
       [
           [0.2, 0.1, 0.7, 0.4],
           [0.9, 0.0, 0.5, 0.3],
           [0.4, 0.4, 0.2, 0.8],
       ],
       dtype=torch.float32,
   )
   targets = torch.tensor(
       [
           [0.5, 0.1],
           [0.8, 0.2],
           [0.3, 0.6],
       ],
       dtype=torch.float32,
   )
    
   model = SensorRegressor()
   loss_fn = nn.MSELoss()
   optimizer = torch.optim.SGD(model.parameters(), lr=0.05)

4. Append the gradient-training block after the optimizer definition.

   predictions = model(features)
   loss = loss_fn(predictions, targets)
   optimizer.zero_grad(set_to_none=True)
   loss.backward()
    
   gradient_norm = torch.linalg.vector_norm(
       torch.stack(
           [
               parameter.grad.detach().norm()
               for parameter in model.parameters()
               if parameter.grad is not None
           ]
       ),
       2,
   )
   output_weight_before = model.output.weight.detach().clone()
   optimizer.step()
   output_weight_changed = not torch.equal(
       model.output.weight.detach(), output_weight_before
   )

   gradient_norm proves that backpropagation reached registered parameters, while output_weight_changed records the optimizer's effect.
   Related: How to zero gradients in PyTorch

5. Append the result checks after the optimizer update.

   registered_layers = list(dict(model.named_children()))
   parameter_names = [name for name, _ in model.named_parameters()]
    
   assert registered_layers == ["input", "activation", "output"]
   assert parameter_names == [
       "input.weight",
       "input.bias",
       "output.weight",
       "output.bias",
   ]
   assert predictions.shape == (3, 2)
   assert gradient_norm.item() > 0
   assert output_weight_changed
    
   print(f"registered_layers={registered_layers}")
   print(f"parameter_names={parameter_names}")
   print(f"output_shape={tuple(predictions.shape)}")
   print(f"gradient_norm={gradient_norm.item():.6f}")
   print(f"output_weight_changed={output_weight_changed}")

6. Run the completed smoke script.

   $ python custom_module_demo.py
   registered_layers=['input', 'activation', 'output']
   parameter_names=['input.weight', 'input.bias', 'output.weight', 'output.bias']
   output_shape=(3, 2)
   gradient_norm=1.048025
   output_weight_changed=True