What is the vanishing/exploding gradient problem and how do you detect and fix it in PyTorch?

import torch
import torch.nn as nn

model     = nn.LSTM(input_size=10, hidden_size=128, num_layers=3, batch_first=True)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

x = torch.randn(32, 20, 10)
output, _ = model(x)
loss = output.sum()

optimizer.zero_grad()
loss.backward()

# ── Detect: monitor gradient norms
total_norm = 0.0
for p in model.parameters():
    if p.grad is not None:
        total_norm += p.grad.data.norm(2).item() ** 2
total_norm = total_norm ** 0.5
print(f"Gradient norm: {total_norm:.4f}")
# Very small (~1e-6) → vanishing; very large (~1e3+) → exploding

# ── Fix 1: Gradient clipping — caps the gradient norm before the step
nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
optimizer.step()

# ── Fix 2: Better weight initialisation (He init for ReLU networks)
def init_weights(m):
    if isinstance(m, nn.Linear):
        nn.init.kaiming_uniform_(m.weight, nonlinearity="relu")
        nn.init.zeros_(m.bias)

mlp = nn.Sequential(nn.Linear(64, 64), nn.ReLU(), nn.Linear(64, 64))
mlp.apply(init_weights)

# ── Fix 3: Batch Normalization — stabilises layer input distributions
class StableNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(64, 64), nn.BatchNorm1d(64), nn.ReLU(),
            nn.Linear(64, 64), nn.BatchNorm1d(64), nn.ReLU(),
        )
    def forward(self, x):
        return self.net(x)

# ── Fix 4: Residual / skip connections — gradient highway
class ResidualBlock(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.net = nn.Sequential(nn.Linear(dim, dim), nn.ReLU(), nn.Linear(dim, dim))
    def forward(self, x):
        return x + self.net(x)   # gradient flows through x directly