Issues with Classification problem

[AlessandroPela]

I am currently at In[25] on chapter 02 https://www.learnpytorch.io/02_pytorch_classification,
model_1 and model_0 both works as in the tutorial, but after I plot the predictions with plot_decision_boundary from https://raw.githubusercontent.com/mrdbourke/pytorch-deep-learning/main/helper_functions.py I tried to the model_1 shell one more time, in order to train 1000 more epochs. I received the following error message:

RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cpu and cuda:0! (when checking argument for argument mat1 in method wrapper_CUDA_addmm)

I saw that without plotting the predictions in the middle of the 2 training everything works, so I imagine that when I un plot_decision_boundary X_train is not any more usable for training.
How can I solve?

[Perian-Yan]

Hi,
I have the same problem as you.
It seems that in the plot_decision_boundary function, the model is moved to cpu and didn't move back. Adding model_0.to(device) in cell In[25] resolves the problem.

[Perian-Yan]

Sorry, its cell In[20], but anyway, just move both the data and model to gpu