Correctly cast num_train_epochs to int (huggingface#11379)

Iwontbecreative · Jul 15, 2021 · 13c6e32 · 13c6e32
1 parent 4958429
commit 13c6e32
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/examples/tensorflow/text-classification/run_text_classification.py b/examples/tensorflow/text-classification/run_text_classification.py
@@ -492,7 +492,10 @@ def preprocess_function(examples):
 
         callbacks = [SavePretrainedCallback(output_dir=training_args.output_dir)]
         model.fit(
-            training_dataset, validation_data=eval_dataset, epochs=training_args.num_train_epochs, callbacks=callbacks
+            training_dataset,
+            validation_data=eval_dataset,
+            epochs=int(training_args.num_train_epochs),
+            callbacks=callbacks,
         )
     elif "validation" in datasets:
         # If there's a validation dataset but no training set, just evaluate the metrics