How would you implement and train a simple feedforward neural network in PyTorch from scratch, without using nn.Sequential?

This question tests whether you understand the full PyTorch workflow: defining a custom nn.Module, implementing forward, and running the standard train loop. It is a common practical screen in ML engineering interviews.

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset

# ─── 1. Define the model ────────────────────────────────────────────
class FeedForwardNet(nn.Module):
    def __init__(self, in_dim: int, hidden_dim: int, out_dim: int,
                 dropout: float = 0.1):
        super().__init__()
        self.fc1     = nn.Linear(in_dim, hidden_dim)
        self.bn1     = nn.BatchNorm1d(hidden_dim)
        self.relu    = nn.ReLU()
        self.drop    = nn.Dropout(dropout)
        self.fc2     = nn.Linear(hidden_dim, out_dim)
        self._init_weights()

    def _init_weights(self):
        nn.init.kaiming_uniform_(self.fc1.weight, nonlinearity='relu')
        nn.init.zeros_(self.fc1.bias)
        nn.init.xavier_uniform_(self.fc2.weight)
        nn.init.zeros_(self.fc2.bias)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.relu(self.bn1(self.fc1(x)))
        x = self.drop(x)
        return self.fc2(x)

# ─── 2. Create data ──────────────────────────────────────────────────
torch.manual_seed(42)
X = torch.randn(1000, 20)
y = (X[:, 0] + X[:, 1] > 0).long()  # binary label
ds     = TensorDataset(X, y)
loader = DataLoader(ds, batch_size=64, shuffle=True)

# ─── 3. Instantiate model, loss, optimizer ───────────────────────────
device    = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model     = FeedForwardNet(20, 64, 2, dropout=0.1).to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.AdamW(model.parameters(), lr=1e-3, weight_decay=1e-4)

# ─── 4. Training loop ─────────────────────────────────────────────────
for epoch in range(30):
    model.train()
    epoch_loss = 0.0
    for X_b, y_b in loader:
        X_b, y_b = X_b.to(device), y_b.to(device)
        optimizer.zero_grad(set_to_none=True)
        loss = criterion(model(X_b), y_b)
        loss.backward()
        optimizer.step()
        epoch_loss += loss.item()
    if epoch % 5 == 0:
        print(f'Epoch {epoch:3d}: loss={epoch_loss / len(loader):.4f}')

Key interview checkpoints: (1) subclass nn.Module and call super().__init__(); (2) define all layers as attributes in __init__; (3) implement forward; (4) follow the zero-grad → forward → loss → backward → step order; (5) call model.train() before training and model.eval() before evaluation.