diff --git a/unsloth/chat_templates.py b/unsloth/chat_templates.py
index 1b59e379a7..70599b8b41 100644
--- a/unsloth/chat_templates.py
+++ b/unsloth/chat_templates.py
@@ -1842,10 +1842,10 @@ def train_on_responses_only(
         return { "labels" : all_labels }
     pass
 
-    if hasattr(trainer, "train_dataset"):
+    if hasattr(trainer, "train_dataset") and trainer.train_dataset is not None:
         trainer.train_dataset = trainer.train_dataset.map(_train_on_responses_only, batched = True)
-    if hasattr(trainer, "eval_dataset"):
-        trainer.eval_dataset  = trainer.eval_dataset .map(_train_on_responses_only, batched = True)
+    if hasattr(trainer, "eval_dataset") and trainer.eval_dataset is not None:
+        trainer.eval_dataset = trainer.eval_dataset.map(_train_on_responses_only, batched = True)
     return trainer
 pass