use numel instead of dim in engine logging

Cadene · Jul 17, 2020 · 05b7ce5 · 05b7ce5
1 parent 728c283
commit 05b7ce5
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/bootstrap/engines/engine.py b/bootstrap/engines/engine.py
@@ -200,8 +200,8 @@ def train_epoch(self, model, dataset, optimizer, epoch, mode='train'):
 
             for key, value in out.items():
                 if torch.is_tensor(value):
-                    if value.dim() <= 1:
-                        value = value.item()  # get number from a torch scalar
+                    if value.numel() <= 1:
+                        value = value.item() # get number from a torch scalar
                     else:
                         continue
                 if isinstance(value, (list, dict, tuple)):