Update train_reader.py

fix 0 train loss logging
facebookresearch · Jan 15, 2022 · 25ed1ff · 25ed1ff
1 parent baf533c
commit 25ed1ff
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/train_reader.py b/train_reader.py
@@ -78,11 +78,11 @@ def train(model, optimizer, scheduler, step, train_dataset, eval_dataset, opt, c
                     log += f"train: {curr_loss/opt.eval_freq:.3f} |"
                     log += f"evaluation: {100*dev_em:.2f}EM |"
                     log += f"lr: {scheduler.get_last_lr()[0]:.5f}"
-                    logger.info(log)
-                    curr_loss = 0
+                    logger.info(log)    
                     if tb_logger is not None:
                         tb_logger.add_scalar("Evaluation", dev_em, step)
                         tb_logger.add_scalar("Training", curr_loss / (opt.eval_freq), step)
+                    curr_loss = 0.
 
             if opt.is_main and step % opt.save_freq == 0:
                 src.util.save(model, optimizer, scheduler, step, best_dev_em,