Solved by albanD in post #12 
              
              
                This is indeed unrelated. 
If you enable anomaly mode, you will see that the problem is that some of the params are saved for backward but modified inplace. The fix is to make sure they are not: 
for i in range(epochs):
    model.train()
    train_loss = 0
    params = dict(model.named_parameters())…