[Examples] Check key exists in datasets first (#11503)

2021-05-10 03:42:38 +08:00
parent ba0d50f214
commit f2ffcaf49f
3 changed files with 3 additions and 3 deletions
--- a/examples/pytorch/summarization/run_summarization.py
+++ b/examples/pytorch/summarization/run_summarization.py
@@ -422,9 +422,9 @@ def main():
        return model_inputs

    if training_args.do_train:
-        train_dataset = datasets["train"]
        if "train" not in datasets:
            raise ValueError("--do_train requires a train dataset")
+        train_dataset = datasets["train"]
        if data_args.max_train_samples is not None:
            train_dataset = train_dataset.select(range(data_args.max_train_samples))
        train_dataset = train_dataset.map(