diff --git a/requirements.txt b/requirements.txt
index 06999de6..d39d3042 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -19,6 +19,7 @@ sentencepiece
 tiktoken
 modelscope>=1.14.0
 hf-transfer
+safetensors<=0.5.3
 # python
 fire
 omegaconf
diff --git a/src/llamafactory/hparams/data_args.py b/src/llamafactory/hparams/data_args.py
index 4c8351f8..90036706 100644
--- a/src/llamafactory/hparams/data_args.py
+++ b/src/llamafactory/hparams/data_args.py
@@ -35,7 +35,7 @@ class DataArguments:
         default=None,
         metadata={"help": "The name of dataset(s) to use for evaluation. Use commas to separate multiple datasets."},
     )
-    dataset_dir: Union[str, dict] = field(
+    dataset_dir: str = field(
         default="data",
         metadata={"help": "Path to the folder containing the datasets."},
     )