S01: Putting everything together - plot_predictions needs to be device agnostic

[sengopal]

While going through "Section 01: Python workflow" - "Putting Everything together" (Part 4), the plot_predictions does not work with the code by only moving test_preds to cpu. The X_train, X_test and y_train, y_test also have to be moved to cpu for plot_predictions to work. Is my understanding correct? It fails on collab if those objects are also not moved to cpu

[mrdbourke]

Yes, that's correct if you're making predictions and would like to plot them, they should be on the cpu (using .cpu() or something similar), otherwise you will get device errors as matplotlib is not compatible with GPU.

See this section for more: https://www.learnpytorch.io/01_pytorch_workflow/#64-making-predictions

And if you're still getting issues, please post your code/let me know what's happening.

[sengopal]

Right :) But the code in the video and also the notebook moves only the test_preds and not the other variables to cpu - :)

[sengopal]

Please check the cell before the section "6.5 Saving & loading a trained model". Only plot_predictions(predictions=y_preds.cpu()), but not the other variables are moved to cpu()

[haopengcu]

I happened to notice this issue. The function 'plot_predictions' has X_train, y_train, X_test, y_test as default inputs. Before we train the simple linear model and move data to gpu, we have already called this function once to check the prediction of the initialized parameters. There this plot function somehow MEMORIZED those default data on the cpu. Next when we plot the predictions after model training, the implicit default inputs will still load the previously called data on the cpu and run through even the train/test data have been moved to gpu. To erase that memory, we can rerun the cell defining the plot function, then we will meet with the error of inconsistent device raised if simply call plot_predictions(predictions=y_preds.cpu()) . There we have to replace the default inputs with explicit X_train.cpu() etc to fix the error.