Hi,
Yes it will work as this just change how you compute the gradients, nothing else. An optimizer like adam is agnostic to the way you obtained your gradients.
In your code you want to do: loss_sum += loss.item() to make sure you do not keep track of the history of all your losses. .item() (or you could use .detach()) will break the graph and thus allow it to be freed from one iteration of the loop to the next.