Fix #1983 by moving input and model to device (#2249)

heloa-net · Svetlana Karslioglu · web-flow · commit 7939159f2d83 · 2023-04-24T10:57:15.000-07:00
Fixes #1983 Moved input to device on line `234` instead of `232` for readability. `y` on line `232` can't be moved to device due to being an `int` I've found out about this issue while following the tutorial from Quickstart and reaching "Optimization model parameters". I was iterating on the code started on Quickstart and the model had been sent to the device. [Optimization tutorial] (https://github.com/pytorch/tutorials/blob/main/beginner_source/basics/optimization_tutorial.py) doesn't move the model or input to the device and this is when I got the error described in #1983. In light of this, should optimization_tutorial.py also be updated? Test plan: * Tested in Colab using GPU, [link](https://colab.research.google.com/drive/1X7BJHd42FN9N_8DiXd7h4ScplN2-OmLo?usp=sharing) Co-authored-by: Svetlana Karslioglu <svekars@fb.com>
diff --git a/beginner_source/basics/quickstart_tutorial.py b/beginner_source/basics/quickstart_tutorial.py
@@ -215,7 +215,7 @@ def test(dataloader, model, loss_fn):
 # The process for loading a model includes re-creating the model structure and loading
 # the state dictionary into it.
 
-model = NeuralNetwork()
+model = NeuralNetwork().to(device)
 model.load_state_dict(torch.load("model.pth"))
 
 #############################################################
@@ -237,6 +237,7 @@ def test(dataloader, model, loss_fn):
 model.eval()
 x, y = test_data[0][0], test_data[0][1]
 with torch.no_grad():
+    x = x.to(device)
     pred = model(x)
     predicted, actual = classes[pred[0].argmax(0)], classes[y]
     print(f'Predicted: "{predicted}", Actual: "{actual}"')
