can you print your device and data inside the training. Just to. confirm weather the data itself is getting transferred to the GPU. When you print the data t will show tensor ending with cuda device id