fix possibly OOM error

Former-commit-id: e6126244c161dc87b0d4d45b8976c02fc9933545
This commit is contained in:
hiyouga 2023-06-01 23:54:44 +08:00
parent 3f3b475412
commit 7e1be4c21a

View File

@ -477,13 +477,13 @@ def preprocess_data(
desc="Running tokenizer on dataset"
)
if stage == "pt":
print_unsupervised_dataset_example(dataset[0])
elif stage == "sft":
print_supervised_dataset_example(dataset[0])
elif stage == "rm":
print_pairwise_dataset_example(dataset[0])
elif stage == "ppo":
print_unsupervised_dataset_example(dataset[0])
if stage == "pt":
print_unsupervised_dataset_example(dataset[0])
elif stage == "sft":
print_supervised_dataset_example(dataset[0])
elif stage == "rm":
print_pairwise_dataset_example(dataset[0])
elif stage == "ppo":
print_unsupervised_dataset_example(dataset[0])
return dataset
return dataset