fix reloading model for evaluation in examples

2018-12-13 14:48:12 +01:00
parent 0f544625f4
commit 087798b7fa
4 changed files with 65 additions and 17 deletions
--- a/examples/run_squad.py
+++ b/examples/run_squad.py
@@ -923,6 +923,9 @@ def main():
    # Load a trained model that you have fine-tuned
    model_state_dict = torch.load(output_model_file)
    model = BertForQuestionAnswering.from_pretrained(args.bert_model, state_dict=model_state_dict)
+    if args.fp16:
+        model.half()
+    model.to(device)

    if args.do_predict and (args.local_rank == -1 or torch.distributed.get_rank() == 0):
        eval_examples = read_squad_examples(