02_ Classification, tensors not in same device

[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?

[mrdbourke]

Hi @AlessandroPela ,

In helper_functions.py, the plot_decision_boundary() function moves the target tensors to the 'cpu', see:

def plot_decision_boundary(model: torch.nn.Module, X: torch.Tensor, y: torch.Tensor):
    """Plots decision boundaries of model predicting on X in comparison to y.

    Source - https://madewithml.com/courses/foundations/neural-networks/ (with modifications)
    """
    # Put everything to CPU (works better with NumPy + Matplotlib)
    model.to("cpu")
    X, y = X.to("cpu"), y.to("cpu")

    # Setup prediction boundaries and grid
    x_min, x_max = X[:, 0].min() - 0.1, X[:, 0].max() + 0.1
    y_min, y_max = X[:, 1].min() - 0.1, X[:, 1].max() + 0.1
    xx, yy = np.meshgrid(np.linspace(x_min, x_max, 101), np.linspace(y_min, y_max, 101))

    # Make features
    X_to_pred_on = torch.from_numpy(np.column_stack((xx.ravel(), yy.ravel()))).float()

    # Make predictions
    model.eval()
    with torch.inference_mode():
        y_logits = model(X_to_pred_on)

    # Test for multi-class or binary and adjust logits to prediction labels
    if len(torch.unique(y)) > 2:
        y_pred = torch.softmax(y_logits, dim=1).argmax(dim=1)  # mutli-class
    else:
        y_pred = torch.round(torch.sigmoid(y_logits))  # binary

    # Reshape preds and plot
    y_pred = y_pred.reshape(xx.shape).detach().numpy()
    plt.contourf(xx, yy, y_pred, cmap=plt.cm.RdYlBu, alpha=0.7)
    plt.scatter(X[:, 0], X[:, 1], c=y, s=40, cmap=plt.cm.RdYlBu)
    plt.xlim(xx.min(), xx.max())
    plt.ylim(yy.min(), yy.max())

So in order to train again, you'll need to make sure all of your tensors/models are on the same device.

You can do this via X_train.to(device) as well as model.to(device) and repeat the same for each of the other tensor/model objects you'd like on the target device.

When training/testing, PyTorch expects all objects (tensors and models) on the same device (e.g. all on GPU or all on CPU).

[mrdbourke]

Hi @AlessandroPela ,

Did this answer your question?

If so, please click the "Mark as answer" button above.

Otherwise, let me know if you have any follow ups.