From c7ed02955347ce9e2b1c87fd7d5e6a0da8fc2112 Mon Sep 17 00:00:00 2001 From: Benjamin Fineran Date: Mon, 26 Apr 2021 13:25:48 -0400 Subject: [PATCH] Update run_qa.py --- integrations/transformers/run_qa.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/integrations/transformers/run_qa.py b/integrations/transformers/run_qa.py index 8aab2bc66de..b6f683bc13d 100644 --- a/integrations/transformers/run_qa.py +++ b/integrations/transformers/run_qa.py @@ -718,7 +718,7 @@ def prepare_validation_features(examples): optim = load_optimizer(model, TrainingArguments) steps_per_epoch = math.ceil(len(datasets["train"]) / (training_args.per_device_train_batch_size*training_args._n_gpu)) manager = ScheduledModifierManager.from_yaml(data_args.nm_prune_config) - training_args.num_train_epochs = float(manager.modifiers[0].end_epoch) + training_args.num_train_epochs = float(manager.max_epochs) optim = ScheduledOptimizer(optim, model, manager, steps_per_epoch=steps_per_epoch, loggers=None) #################################################################################### # End SparseML Integration