[Flax] token-classification model steps enumerate start from 1 (#14547)

* step start from 1 * Updated cur_step calcualtion
huggingface · Nov 29, 2021 · 2bd950c · 2bd950c
1 parent cea17ac
commit 2bd950c
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/examples/flax/token-classification/run_flax_ner.py b/examples/flax/token-classification/run_flax_ner.py
@@ -598,7 +598,7 @@ def compute_metrics():
             state, train_metric, dropout_rngs = p_train_step(state, batch, dropout_rngs)
             train_metrics.append(train_metric)
 
-            cur_step = epoch * step_per_epoch + step
+            cur_step = (epoch * step_per_epoch) + (step + 1)
 
             if cur_step % training_args.logging_steps == 0 and cur_step > 0:
                 # Save metrics