John1231983:

In your code you want to do:  loss_sum += loss.item()


Thanks. So you meant
optimizer.zero_grad()
loss_sum = 0
for  i in range(iter_size):
    loss = criterion(output, target_var) / iter_size
    loss.backward()
    loss_sum += loss.item() # Change here
optimizer.step()