What is nn.Module and how do you build a custom neural network in PyTorch?

import torch
import torch.nn as nn

class MLP(nn.Module):
    def __init__(self, in_features: int, hidden: int, out_features: int):
        super().__init__()   # MUST call this first!

        # Layers defined as attributes are auto-registered as sub-modules
        self.fc1  = nn.Linear(in_features, hidden)
        self.relu = nn.ReLU()
        self.drop = nn.Dropout(p=0.3)
        self.fc2  = nn.Linear(hidden, out_features)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Define the forward computation."""
        x = self.fc1(x)
        x = self.relu(x)
        x = self.drop(x)
        x = self.fc2(x)
        return x

# Instantiate and inspect
model = MLP(in_features=784, hidden=256, out_features=10)

# Forward pass — calls forward() via __call__
x = torch.randn(32, 784)   # batch of 32
out = model(x)              # shape (32, 10)

# Parameter inspection
for name, param in model.named_parameters():
    print(name, param.shape, param.requires_grad)
# fc1.weight  torch.Size([256, 784])  True
# fc1.bias    torch.Size([256])       True
# fc2.weight  torch.Size([10, 256])   True
# fc2.bias    torch.Size([10])        True

total = sum(p.numel() for p in model.parameters())
print(f"Total parameters: {total:,}")