make sure model on target device

mrdbourke · web-flow · commit defa4faa2339 · 2022-07-27T16:06:34.000+10:00
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)):

Original file line number	Diff line number	Diff line change
`@@ -160,6 +160,9 @@ def train(model: torch.nn.Module,`
`160`	`160`	`"test_loss": [],`
`161`	`161`	`"test_acc": []`
`162`	`162`	`}`
	`163`	`+`
	`164`	`+ # Make sure model on target device`
	`165`	`+ model.to(device)`
`163`	`166`
`164`	`167`	`# Loop through training and testing steps for a number of epochs`
`165`	`168`	`for epoch in tqdm(range(epochs)):`