Merge branch 'update-exmaples-versions' of https://github.com/adap/fl…

…ower into update-exmaples-versions
adap · Jan 31, 2025 · d3e2f1b · d3e2f1b
2 parents ab2c922 + 8388f01
commit d3e2f1b
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/examples/embedded-devices/embeddedexample/task.py b/examples/embedded-devices/embeddedexample/task.py
@@ -84,6 +84,8 @@ def train(net, trainloader, valloader, epochs, learning_rate, device):
 
 def test(net, testloader, device):
     """Validate the model on the test set."""
+    net.to(device)  # move model to GPU if available
+    net.eval()
     criterion = torch.nn.CrossEntropyLoss()
     correct, loss = 0, 0.0
     with torch.no_grad():