Fix order of loss resets. (#259)

* Fix order of loss resets. * Formatting. --------- Co-authored-by: George Smyrnis <[email protected]>
mlfoundations · May 8, 2024 · 6321543 · 6321543
1 parent 87955f5
commit 6321543
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 6 deletions.
diff --git a/open_lm/train.py b/open_lm/train.py
@@ -326,11 +326,6 @@ def train_one_epoch(
                 # resetting batch / data time meters per log window
                 batch_time_m.reset()
                 data_time_m.reset()
-                # reset all average meters
-                losses_m.reset()
-                if averagers is not None and args.log_avg_model_training_loss:
-                    for k in averagers.avgs_dict.keys():
-                        losses_avg_m[k].reset()
 
                 if math.isnan(losses_m.val):
                     # case where loss goes to nan, we see this sometimes with bad nodes.
@@ -339,6 +334,12 @@ def train_one_epoch(
                     # training on restarts.
                     return False, step
 
+                # reset all average meters
+                losses_m.reset()
+                if averagers is not None and args.log_avg_model_training_loss:
+                    for k in averagers.avgs_dict.keys():
+                        losses_avg_m[k].reset()
+
     # end for
     if tb_writer is not None:
         tb_writer.flush()

diff --git a/setup.py b/setup.py
@@ -19,7 +19,7 @@ def _read_reqs(relpath):
 
 setuptools.setup(
     name="open_lm",
-    version="0.0.33",
+    version="0.0.34",
     author=[
         "Suchin Gururangan*",
         "Mitchell Wortsman*",