From 261739609473a63bd3707c46221cf88990400a11 Mon Sep 17 00:00:00 2001 From: Matt Date: Thu, 22 Apr 2021 13:49:59 +0100 Subject: [PATCH] Correctly cast num_train_epochs to int (#11379) --- .../text-classification/run_text_classification.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/examples/tensorflow/text-classification/run_text_classification.py b/examples/tensorflow/text-classification/run_text_classification.py index 861a9ccc3c..3c9e260097 100644 --- a/examples/tensorflow/text-classification/run_text_classification.py +++ b/examples/tensorflow/text-classification/run_text_classification.py @@ -492,7 +492,10 @@ def main(): 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