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.