From 9c88eb511d406e7712f6962ae0c3897cfd112fad Mon Sep 17 00:00:00 2001
From: Csaba Kecskemeti <csaba.kecskemeti@gmail.com>
Date: Sat, 24 Feb 2024 10:12:13 -0800
Subject: [PATCH] Update pytorch-ddp-accelerate-transformers.md

As the model is on 'cuda', the test data and labels also should be on the same device for evaluation. I guess this is just a copy paste issue
---
 pytorch-ddp-accelerate-transformers.md | 1 +
 1 file changed, 1 insertion(+)

diff --git a/pytorch-ddp-accelerate-transformers.md b/pytorch-ddp-accelerate-transformers.md
index d679804475..8e4d46f235 100644
--- a/pytorch-ddp-accelerate-transformers.md
+++ b/pytorch-ddp-accelerate-transformers.md
@@ -101,6 +101,7 @@ model.eval()
 correct = 0
 with torch.no_grad():
     for data, target in test_loader:
+        data, target = data.to(device), target.to(device)
         output = model(data)
         pred = output.argmax(dim=1, keepdim=True)
         correct += pred.eq(target.view_as(pred)).sum().item()