Merge pull request #8 from junweima/fix

fix 0 percent training accuracy
facebookresearch · Sep 8, 2018 · eaff31a · eaff31a
2 parents 7af1916 + 044c560
commit eaff31a
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()