Fix it to work with BART (huggingface#6756)

Zigur · Oct 26, 2020 · d00833b · d00833b
1 parent bcc65f7
commit d00833b
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/examples/question-answering/run_squad.py b/examples/question-answering/run_squad.py
@@ -187,7 +187,7 @@ def train(args, train_dataset, model, tokenizer):
                 "end_positions": batch[4],
             }
 
-            if args.model_type in ["xlm", "roberta", "distilbert", "camembert"]:
+            if args.model_type in ["xlm", "roberta", "distilbert", "camembert", "bart"]:
                 del inputs["token_type_ids"]
 
             if args.model_type in ["xlnet", "xlm"]:
@@ -300,7 +300,7 @@ def evaluate(args, model, tokenizer, prefix=""):
                 "token_type_ids": batch[2],
             }
 
-            if args.model_type in ["xlm", "roberta", "distilbert", "camembert"]:
+            if args.model_type in ["xlm", "roberta", "distilbert", "camembert", "bart"]:
                 del inputs["token_type_ids"]
 
             feature_indices = batch[3]