From e6126244c161dc87b0d4d45b8976c02fc9933545 Mon Sep 17 00:00:00 2001 From: hiyouga Date: Thu, 1 Jun 2023 23:54:44 +0800 Subject: [PATCH] fix possibly OOM error --- src/utils/common.py | 20 ++++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/src/utils/common.py b/src/utils/common.py index 9137e54..2652328 100644 --- a/src/utils/common.py +++ b/src/utils/common.py @@ -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]) - - return 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]) + + return dataset