From be3b9bcf4db4f7e942c7f71eb1d7de3a8d476ad0 Mon Sep 17 00:00:00 2001 From: Jade Abbott Date: Thu, 3 Jan 2019 09:02:33 +0200 Subject: [PATCH] Allow one to use the pretrained model in evaluation when do_train is not selected --- examples/run_classifier.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/examples/run_classifier.py b/examples/run_classifier.py index adf81f4e28..9236c6a252 100644 --- a/examples/run_classifier.py +++ b/examples/run_classifier.py @@ -430,8 +430,8 @@ def main(): if not args.do_train and not args.do_eval: raise ValueError("At least one of `do_train` or `do_eval` must be True.") - - if os.path.exists(args.output_dir) and os.listdir(args.output_dir): + + if os.path.exists(args.output_dir) and os.listdir(args.output_dir) and args.do_train: raise ValueError("Output directory ({}) already exists and is not empty.".format(args.output_dir)) os.makedirs(args.output_dir, exist_ok=True) @@ -554,7 +554,8 @@ def main(): # Save a trained model model_to_save = model.module if hasattr(model, 'module') else model # Only save the model it-self output_model_file = os.path.join(args.output_dir, "pytorch_model.bin") - torch.save(model_to_save.state_dict(), output_model_file) + if args.do_train: + torch.save(model_to_save.state_dict(), output_model_file) # Load a trained model that you have fine-tuned model_state_dict = torch.load(output_model_file)