Update run_qa.py (#195)

Author: bfineran

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