Solved by ngimel in post #6 
              
              
                @ptrblck is right, it’s coming from accumulating losses, a simple repro is 
a=torch.tensor(2.)
b=torch.tensor(2., device="cuda")
a += b

The workaround is 
a = a + b

(a would be on the gpu in this case) 
I’ve opened an upstream issue to track