What is Batch Normalization in PyTorch and how does it differ from Layer Normalization?

import torch
import torch.nn as nn

# ── BatchNorm — for feedforward / CNN models
class BNModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(20, 64)
        self.bn1 = nn.BatchNorm1d(64)   # 64 features
        self.fc2 = nn.Linear(64, 10)

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

# BatchNorm behaves differently in train vs eval mode!
# train: normalise using batch mean/var, update running stats
# eval:  use accumulated running_mean / running_var
model = BNModel()
model.train()   # must be in train mode during training!

# ── LayerNorm — for transformers and sequence models
class LNModel(nn.Module):
    def __init__(self, d_model=64):
        super().__init__()
        self.fc1 = nn.Linear(20, d_model)
        self.ln1 = nn.LayerNorm(d_model)   # normalise over last dim
        self.fc2 = nn.Linear(d_model, 10)

    def forward(self, x):
        x = torch.relu(self.ln1(self.fc1(x)))
        return self.fc2(x)

# LayerNorm produces the SAME result at train and eval
ln_model = LNModel()
ln_model.train()
x = torch.randn(8, 20)
out_train = ln_model(x)
ln_model.eval()
out_eval  = ln_model(x)
print(torch.allclose(out_train, out_eval))  # True — LayerNorm is mode-independent!

Common bug: forgetting to call model.train() before training and model.eval() before validation when using BatchNorm — at eval, it uses accumulated running statistics, and if these were never updated (because the model was always in eval mode), predictions will be incorrect.