SimonW:

PyTorch uses dynamic graph. It’s not symbolic expression as tf is. You are backproping via this graph multiple times. Manually compute at each iteration please.


This was the trick! I had exactly the conceptual confusion that you’ve identified. It was also the source of my problems in my original code.
Thanks!