From 37baa3aa254af721728aa33befdc383858cb8ea2 Mon Sep 17 00:00:00 2001 From: Volpeon Date: Sun, 15 Jan 2023 10:38:49 +0100 Subject: Removed unused code, put training callbacks in dataclass --- training/strategy/ti.py | 20 ++++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) (limited to 'training/strategy') 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, + ) -- cgit v1.2.3-70-g09d2