From 8120bb875e8bebcd3988b0a2bb860b3b6c1894bd Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 21 Nov 2024 01:49:46 -0800 Subject: [PATCH] Update trainer.py --- unsloth/trainer.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) 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: