Hi @ptrblck,
Thanks a lot for your feedback and tips for the loss. I check the link you proposed and it helps me a lot to understand this forward-backward process. I also play with the optimizer and learning rate to get a smoother gradient descent and everything works well now.
I also was able to reduce my GPU memory from 22Gb to 6Gb using the solution proposed by @albanD (where he presents the 3 training loops).
BTW, thanks a lot for maintaining this forum. I’m using Pytorch for one year and it helps me a lot to dig into its usage =)