How to create a custom nn.Module in PyTorch

Custom PyTorch modules are the point where loose tensor operations become reusable model code. A class based on torch.nn.Module can hold layers, expose trainable parameters to optimizers, and keep the forward pass in one named object that can be saved, loaded, tested, and reused.

A module registers child layers when they are assigned as attributes in __init__(). The forward() method defines how an input tensor moves through those layers, while calling the module instance, such as model(features), lets PyTorch run the normal module call path around that method.

The smoke workflow uses a small CPU regression model with two nn.Linear layers and a ReLU activation. A forward pass proves the output shape, and one optimizer step proves that gradients reach the registered parameters.

1. Create a smoke script with a module class, sample tensors, and one optimizer step.

   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)
    
       def forward(self, features):
           hidden = self.activation(self.input(features))
           return self.output(hidden)
    
    
   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()
   optimizer = torch.optim.SGD(model.parameters(), lr=0.05)
   loss_fn = nn.MSELoss()
    
   predictions = model(features)
   loss = loss_fn(predictions, targets)
    
   optimizer.zero_grad(set_to_none=True)
   loss.backward()
    
   gradient_norms = [
       parameter.grad.detach().norm()
       for parameter in model.parameters()
       if parameter.grad is not None
   ]
   gradient_norm = torch.linalg.vector_norm(torch.stack(gradient_norms), 2)
    
   output_weight_before = model.output.weight.detach().clone()
   optimizer.step()
   weight_delta = (model.output.weight.detach() - output_weight_before).abs().max()
    
   print(f"torch_version={torch.__version__}")
   print("module=SensorRegressor")
   print(f"parameter_tensors={len(list(model.parameters()))}")
   print(f"trainable_parameters={sum(parameter.numel() for parameter in model.parameters())}")
   print(f"output_shape={tuple(predictions.shape)}")
   print(f"loss={loss.item():.6f}")
   print(f"gradient_norm={gradient_norm.item():.6f}")
   print(f"optimizer_step_changed_weight={bool(weight_delta > 0)}")

   super().init() initializes the base module machinery before child layers are assigned. Layers stored in a plain local variable inside __init__() are not registered as child modules.

2. Run the smoke script.

   $ python custom_module_demo.py
   torch_version=2.12.1+cpu
   module=SensorRegressor
   parameter_tensors=4
   trainable_parameters=58
   output_shape=(3, 2)
   loss=0.311184
   gradient_norm=1.048025
   optimizer_step_changed_weight=True

3. Check the module registration and output lines.

   parameter_tensors=4 and trainable_parameters=58 confirm the two Linear layers are registered. output_shape=(3, 2) confirms the module maps a batch of three four-feature rows to two outputs per row.

4. Copy the module class into the project code.

   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)
    
       def forward(self, features):
           hidden = self.activation(self.input(features))
           return self.output(hidden)

   Keep layer names stable when checkpoints or saved state dictionaries will load weights back into the class.

5. Use the module in a training step.

   model = SensorRegressor(in_features=4, hidden_features=8, out_features=2)
   optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4)
   loss_fn = nn.MSELoss()
    
   for features, targets in train_loader:
       optimizer.zero_grad(set_to_none=True)
       predictions = model(features)
       loss = loss_fn(predictions, targets)
       loss.backward()
       optimizer.step()

   Keep features, targets, and the module parameters on the same device.
   Related: How to zero gradients in PyTorch

6. Call the module instance for inference.

   model.eval()
    
   with torch.inference_mode():
       predictions = model(features)

   Calling model(features) uses the module call path that wraps forward() and honors registered hooks. Direct calls such as model.forward(features) skip that dispatch path.

7. Remove the temporary smoke script.

   $ rm custom_module_demo.py