Can't get the training to be done on gpu

Even though I have explicitly tried moving both the model and the input tensors to GPU I get the following error. And on checking images.device and labels.device , it shows CPU

Help would be much appreciated.

The problem here is on this line The .to() method returns a pointer to the model stored in the specified device. Since you are not assigning this new pointer to any variable you are storing a model on GPU but still using the model stored in CPU. You should change the line to be like this:

model =

Hope that helps. :slight_smile:

1 Like