Update train_imagenet.py

libffcv · Jan 29, 2022 · 607d117 · 607d117
1 parent f134cbf
commit 607d117
Showing 1 changed file with 0 additions and 5 deletions.
diff --git a/train_imagenet.py b/train_imagenet.py
@@ -391,11 +391,6 @@ def train_loop(self, epoch, log_level):
                 iterator.set_description(msg)
             ### Logging end
 
-        if log_level > 0:
-            loss = ch.stack(losses).mean().cpu()
-            assert not ch.isnan(loss), 'Loss is NaN!'
-            return loss.item()
-
     @param('validation.lr_tta')
     def val_loop(self, lr_tta):
         model = self.model