diff --git a/unsloth/trainer.py b/unsloth/trainer.py
index 14fd1631df..012be4b0cb 100644
--- a/unsloth/trainer.py
+++ b/unsloth/trainer.py
@@ -21,7 +21,12 @@ import trl
 import inspect
 from trl import SFTTrainer
 from . import is_bfloat16_supported
-from unsloth_zoo.training_utils import unsloth_train as _unsloth_train
+from unsloth_zoo.training_utils import (
+    unsloth_train as _unsloth_train,
+)
+from unsloth_zoo.vision_utils import (
+    UnslothVisionDataCollator,
+)
 from packaging.version import Version
 import dataclasses
 
@@ -30,6 +35,7 @@ __all__ = [
     "UnslothTrainer",
     "unsloth_train",
     "_patch_trl_trainer",
+    "UnslothVisionDataCollator",
 ]
 
 # Unsloth gradient accumulation fix: