How to transfer a program running on CPU to a program running on gpu
Basically. you would have to push the model and data onto the GPU using
.to(device), so that all operations will take place on the device.
Have a look at the Training a classifier tutorial to see how and where to do it.