diff --git a/scripts/finetune.py b/scripts/finetune.py
index ee055559f..5b7f8f2ab 100644
--- a/scripts/finetune.py
+++ b/scripts/finetune.py
@@ -182,10 +182,6 @@ def train(
         tokenizer, cfg, DEFAULT_DATASET_PREPARED_PATH
     )
 
-    if prepare_ds_only:
-        logging.info("Finished preparing dataset. Exiting...")
-        return
-
     if cfg.debug:
         logging.info("check_dataset_labels...")
         check_dataset_labels(
@@ -195,6 +191,10 @@ def train(
             tokenizer,
         )
 
+    if prepare_ds_only:
+        logging.info("Finished preparing dataset. Exiting...")
+        return
+
     trainer = setup_trainer(cfg, train_dataset, eval_dataset, model, tokenizer)
 
     model.config.use_cache = False