Por que chamar optimizer.zero_grad() no início do ciclo de treino?

No PyTorch, durante a retropropagação, os gradientes dos parâmetros acumulam-se, ou seja, cada chamada a loss.backward() adiciona os gradientes aos já existentes nos tensores .grad. Se não os zerarmos antes de um novo passo de treino, os gradientes serão somados, o que levará a uma atualização incorreta dos pesos.

Por isso, no início de cada ciclo de treino, chama-se optimizer.zero_grad() para limpar os gradientes antigos e calcular corretamente os novos para o lote atual.

Exemplo:

for data, target in dataloader:
    optimizer.zero_grad()  # zera os gradientes
    output = model(data)
    loss = loss_fn(output, target)
    loss.backward()       # calcula os novos gradientes
    optimizer.step()      # atualiza os parâmetros