Hi Koramajin!



 Koramajin:

But when I move all tensors to cpu, the result is correct.


Could this be the TF32 “bug?”
What specific gpu are you using and does the issue go away if you set
torch.backends.cuda.matmul.allow_tf32 = True?
Edit:
The above should read
torch.backends.cuda.matmul.allow_tf32 = False

True is the default for allow_tf32.  The suggestion is to turn off TF32
which is done by setting allow_tf32 to False.
Best.
K. Frank