fix hiyouga#5228

yuwangnexusera · Sep 5, 2024 · 2d6e64b · 2d6e64b
1 parent 2050320
commit 2d6e64b
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/llamafactory/data/loader.py b/src/llamafactory/data/loader.py
@@ -237,7 +237,7 @@ def get_dataset(
             if data_args.streaming:
                 dataset_module = {k: v.to_iterable_dataset() for k, v in dataset_module.items()}
 
-            return dataset_module
+            return dataset_module, template
 
         if data_args.streaming:
             raise ValueError("Turn off `streaming` when saving dataset to disk.")