make sure model on target device

mrdbourke · Jul 27, 2022 · defa4fa · defa4fa
1 parent 6237b93
commit defa4fa
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/going_modular/going_modular/engine.py b/going_modular/going_modular/engine.py
@@ -160,6 +160,9 @@ def train(model: torch.nn.Module,
                "test_loss": [],
                "test_acc": []
     }
+
+    # Make sure model on target device
+    model.to(device)
 
     # Loop through training and testing steps for a number of epochs
     for epoch in tqdm(range(epochs)):