Skip to content

GitLab

Unverified Commit 4605b2b8 authored by Patrick von Platen's avatar Patrick von Platen Committed by GitHub
Browse files

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

* fix_torch_device_generate_test

* remove @

* up
parent d0f7508a
Hide whitespace changes
Inline Side-by-side
Showing with 2 additions and 2 deletions
examples/flax/language-modeling/run_mlm_flax.py
View file @ 4605b2b8
......@@ -606,7 +606,7 @@ if __name__ == "__main__":
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
......
examples/flax/language-modeling/run_t5_mlm_flax.py
View file @ 4605b2b8
......@@ -722,7 +722,7 @@ if __name__ == "__main__":
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
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment