1 files changed, 10 insertions, 10 deletions

diff --git a/training/strategy/ti.py b/training/strategy/ti.py
index 83dc566..6f8384f 100644
--- a/training/strategy/ti.py
+++ b/training/strategy/ti.py
@@ -15,7 +15,7 @@ from slugify import slugify
 
 from models.clip.tokenizer import MultiCLIPTokenizer
 from training.util import EMAModel
-from training.functional import save_samples
+from training.functional import TrainingCallbacks, save_samples
 
 
 def textual_inversion_strategy(
@@ -153,12 +153,12 @@ def textual_inversion_strategy(
         with ema_context:
             save_samples_(step=step)
 
-    return {
-        "on_prepare": on_prepare,
-        "on_train": on_train,
-        "on_eval": on_eval,
-        "on_after_optimize": on_after_optimize,
-        "on_log": on_log,
-        "on_checkpoint": on_checkpoint,
-        "on_sample": on_sample,
-    }
+    return TrainingCallbacks(
+        on_prepare=on_prepare,
+        on_train=on_train,
+        on_eval=on_eval,
+        on_after_optimize=on_after_optimize,
+        on_log=on_log,
+        on_checkpoint=on_checkpoint,
+        on_sample=on_sample,
+    )