[Flax] Fix another bug in logging steps (huggingface#12516)

* fix_torch_device_generate_test * remove @ * up
gregbugaj · Jul 5, 2021 · 4605b2b · 4605b2b
1 parent d0f7508
commit 4605b2b
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/examples/flax/language-modeling/run_mlm_flax.py b/examples/flax/language-modeling/run_mlm_flax.py
@@ -606,7 +606,7 @@ def eval_step(params, batch):
             state, train_metric, dropout_rngs = p_train_step(state, model_inputs, dropout_rngs)
             train_metrics.append(train_metric)
 
-            cur_step = epoch * num_train_samples + step
+            cur_step = epoch * (num_train_samples // train_batch_size) + step
 
             if cur_step % training_args.logging_steps == 0 and cur_step > 0:
                 # Save metrics

diff --git a/examples/flax/language-modeling/run_t5_mlm_flax.py b/examples/flax/language-modeling/run_t5_mlm_flax.py
@@ -722,7 +722,7 @@ def eval_step(params, batch):
             state, train_metric, dropout_rngs = p_train_step(state, model_inputs, dropout_rngs)
             train_metrics.append(train_metric)
 
-            cur_step = epoch * num_train_samples + step
+            cur_step = epoch * (num_train_samples // train_batch_size) + step
 
             if cur_step % training_args.logging_steps == 0 and cur_step > 0:
                 # Save metrics