polm23:

retain_graph=True


This part looks like a “beep” to me.
If you check out the official docs 7 you will try to avoid retain_graph.
Will you try backward w/o the retain_graph=False and let us know. Also, is your train.py somewhere for preview?