some question about the position of 'optimizer.zero_grad()'

[languandong]

I think the correct way the code the training
is that

    optimizer.zero_grad()
    # Forward pass
    outputs = model(images)
    loss = criterion(outputs, labels)
    
    # Backward and optimize
    loss.backward()
    optimizer.step()

not that

    # Forward pass
    outputs = model(images)
    loss = criterion(outputs, labels)
    
    # Backward and optimize
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

[Vandaci]

any difference?

[silky1708]

@languandong
You can use both, doesn't matter as long as optimizer.zero_grad() is called before loss.backward().
Note that optimizer.zero_grad() zeroes out the gradients in the grad field of the tensors, and loss.backward() compute s the gradients which are then stored in the grad field.

[githraj]

As pointed out by @languandong, the critical factor is the correct sequence in which optimizer.zero_grad() and loss.backward() are called. Both code snippets are valid as long as optimizer.zero_grad() is invoked before loss.backward(). This ensures that the gradients are properly zeroed out and then computed and stored in the appropriate tensors' grad field.

[luyuwuli]

@languandong I think the confusion originates from the misconception that the gradient would be computed and stored during the forward pass. In fact, in the forward pass, only the DAG is constructed. The grad is computed in a lazy mode: it is not computed until explicit loss.backward() is invoked.