Solved by albanD in post #2 
              
              
                Hi, 
The .to() operation on Tensor is not inplace, so you need to replace train_full.to(device) by train_full = train_full.to(device).