Hello @joel @albanD,
I am also encountering the same problem in my code. In my case, MSE loss is nan. I tried clipping gradients by adding below lines
torch.nn.utils.clip_grad_norm_(self.net.parameters(), max_norm=1.0, norm_type=2.0)
optimizer.step()
optimizer.zero_grad()
But I still get nan, would you be able to share what actually you modified?