From 05b7ce5b49b99774eb6c6ed16caa7e2e1008e4d1 Mon Sep 17 00:00:00 2001 From: Corentin Dancette Date: Tue, 30 Jun 2020 23:58:53 +0200 Subject: [PATCH] use numel instead of dim in engine logging --- bootstrap/engines/engine.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/bootstrap/engines/engine.py b/bootstrap/engines/engine.py index de37236..86a74a4 100755 --- 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)):