fix 0 percent training accuracy

facebookresearch · Sep 7, 2018 · 044c560 · 044c560
1 parent 7af1916
commit 044c560
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/train.py b/train.py
@@ -158,8 +158,8 @@ def train(epoch):
         train_loss += loss.data[0]
         _, predicted = torch.max(outputs.data, 1)
         total += targets.size(0)
-        correct += (lam * predicted.eq(targets_a.data).cpu().sum()
-                    + (1 - lam) * predicted.eq(targets_b.data).cpu().sum())
+        correct += (lam * predicted.eq(targets_a.data).cpu().sum().float()
+                    + (1 - lam) * predicted.eq(targets_b.data).cpu().sum().float())
 
         optimizer.zero_grad()
         loss.backward()