apply patil-suraj#86

Author: morioka

run_qg.py

@@ -178,14 +178,16 @@ def main(args_file=None):
         using_tpu=training_args.tpu_num_cores is not None
     )
 
+    # Prediction Loss
+    training_args.prediction_loss_only=True
+
     # Initialize our Trainer
     trainer = Trainer(
         model=model,
         args=training_args,
         train_dataset=train_dataset,
         eval_dataset=valid_dataset,
         data_collator=data_collator,
-        prediction_loss_only=True,
         label_smoothing=model_args.label_smoothing
     )
 
@@ -200,8 +202,8 @@ def main(args_file=None):
         trainer.save_model()
         # For convenience, we also re-save the tokenizer to the same directory,
         # so that you can share your model easily on huggingface.co/models =)
-        if trainer.is_world_master():
-            tokenizer.save_pretrained(training_args.output_dir)
+        if trainer.is_world_process_zero():
+                tokenizer.save_pretrained(training_args.output_dir)
 
     # Evaluation
     results = {}
@@ -233,4 +235,4 @@ def run_qg(args_dict):
     main(args_file="args.json")
 
 if __name__ == "__main__":
-    main()
+    main()