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

https://colab.research.google.com/drive/1Sgok4mHvuFT_YF3yHXWFR6RAZXVuCFRD

Help would be much appreciated.

The problem here is on this line model.to(device). 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 = model.to(device)

Hope that helps. :slight_smile:

1 Like