When the loss will not fall, check the code before the model. Can it drive the loss near zero on one small batch? If not, the bug is in the loop: a missing zero_grad, a wrong target shape, a detached tensor.

Then look inside. A parameter whose gradient is exactly zero is disconnected or dead; a nan or infinite gradient means something overflowed. A forward hook is a function PyTorch calls with every output of a module, so you can watch values without changing the model.

import torch
from torch import nn

model = nn.Sequential(
    nn.Linear(3, 4), nn.ReLU(),
    nn.Linear(4, 1))
seen = {}
def keep(name):
    def hook(module, inputs, output):
        seen[name] = output.shape
    return hook
handles = []
for n, m in model.named_children():
    handles.append(
        m.register_forward_hook(keep(n)))
model(torch.randn(2, 3))
for h in handles:
    h.remove()
print(seen)

register_forward_hook returns a handle. Remove it when you are done, or the hook runs on every call for the life of the model.

Write grad_norms(model) returning a dict from each parameter's name to the norm of its gradient (a float), for parameters that have one. Write nonfinite(model) returning the sorted names of parameters whose gradient contains nan or an infinity. Write layer_outputs(model, x) returning a dict from each direct child's name to its output for x, collected with forward hooks that are removed afterwards, without recording gradients.

The explanation opens here once every check passes.