Fix quickstart-pytorch (#2659)

adap · Nov 29, 2023 · 14c9b2a · 14c9b2a
1 parent ded72a2
commit 14c9b2a
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/examples/quickstart-pytorch/client.py b/examples/quickstart-pytorch/client.py
@@ -60,9 +60,9 @@ def test(net, testloader):
     correct, loss = 0, 0.0
     with torch.no_grad():
         for batch in tqdm(testloader, "Testing"):
-            images = batch["img"]
-            labels = batch["label"]
-            outputs = net(images.to(DEVICE))
+            images = batch["img"].to(DEVICE)
+            labels = batch["label"].to(DEVICE)
+            outputs = net(images)
             loss += criterion(outputs, labels).item()
             correct += (torch.max(outputs.data, 1)[1] == labels).sum().item()
     accuracy = correct / len(testloader.dataset)