[QuestionAnswering] catch extra exception for retry in custom file lo…

…ading (#1669)
neuralmagic · Jul 13, 2023 · a0fcc1f · a0fcc1f
1 parent 6d836ed
commit a0fcc1f
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/sparseml/transformers/question_answering.py b/src/sparseml/transformers/question_answering.py
@@ -922,7 +922,7 @@ def _get_raw_dataset(data_args, cache_dir: Optional[str] = None):
                 field="data",
                 cache_dir=cache_dir,
             )
-        except json.JSONDecodeError:
+        except (json.JSONDecodeError, datasets.builder.DatasetGenerationError):
             # run without `field="data"` - JSONL files will not always have each
             # line nested under a top level field
             raw_datasets = load_dataset(