Зачем вызывать optimizer.zero_grad() в начале train-loop?

В PyTorch при обратном распространении ошибки градиенты параметров накапливаются, то есть каждый вызов loss.backward() добавляет градиенты к уже существующим в .grad тензорах. Если не обнулить их перед новым шагом обучения, то градиенты будут суммироваться, что приведёт к неправильному обновлению весов.

Поэтому в начале каждого цикла обучения вызывают optimizer.zero_grad(), чтобы очистить старые градиенты и корректно вычислить новые для текущего батча.

Пример:

for data, target in dataloader:
    optimizer.zero_grad()  # обнуляем градиенты
    output = model(data)
    loss = loss_fn(output, target)
    loss.backward()       # считаем новые градиенты
    optimizer.step()      # обновляем параметры