diff --git a/unsloth/chat_templates.py b/unsloth/chat_templates.py
index e19bea0771..1b59e379a7 100644
--- a/unsloth/chat_templates.py
+++ b/unsloth/chat_templates.py
@@ -1841,7 +1841,11 @@ def train_on_responses_only(
         pass
         return { "labels" : all_labels }
     pass
-    trainer.train_dataset = trainer.train_dataset.map(_train_on_responses_only, batched = True)
+
+    if hasattr(trainer, "train_dataset"):
+        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)
     return trainer
 pass