Merge pull request Diego999#57 from AhMeDxHaMiDo/master

Hey, I fixed an IndexError inside train.py that is faced by many people in the issues section.
xtttf · Aug 15, 2021 · 92ad7af · 92ad7af
2 parents b7b6886 + e469a6e
commit 92ad7af
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/train.py b/train.py
@@ -108,8 +108,8 @@ def compute_test():
     loss_test = F.nll_loss(output[idx_test], labels[idx_test])
     acc_test = accuracy(output[idx_test], labels[idx_test])
     print("Test set results:",
-          "loss= {:.4f}".format(loss_test.data[0]),
-          "accuracy= {:.4f}".format(acc_test.data[0]))
+          "loss= {:.4f}".format(loss_test.data.item()),
+          "accuracy= {:.4f}".format(acc_test.data.item()))
 
 # Train model
 t_total = time.time()