Quick fix to TF summarization example (#14401)

huggingface · Nov 15, 2021 · 267867e · 267867e
1 parent 29dfb2d
commit 267867e
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/examples/tensorflow/summarization/run_summarization.py b/examples/tensorflow/summarization/run_summarization.py
@@ -577,7 +577,7 @@ def postprocess_text(preds, labels):
 
         # region Optimizer, loss and LR scheduling
         # Scheduler and math around the number of training steps.
-        num_update_steps_per_epoch = len(train_dataset) // training_args.per_device_train_batch_size
+        num_update_steps_per_epoch = len(train_dataset) // total_train_batch_size
         num_train_steps = training_args.num_train_epochs * num_update_steps_per_epoch
         optimizer, lr_schedule = create_optimizer(
             init_lr=training_args.learning_rate, num_train_steps=num_train_steps, num_warmup_steps=0