1 files changed, 2 insertions, 7 deletions

diff --git a/training/functional.py b/training/functional.py
index a5b339d..ee73ab2 100644
--- a/training/functional.py
+++ b/training/functional.py
@@ -34,7 +34,6 @@ def const(result=None):
 
 @dataclass
 class TrainingCallbacks():
-    on_prepare: Callable[[], None] = const()
     on_accum_model: Callable[[], torch.nn.Module] = const(None)
     on_log: Callable[[], dict[str, Any]] = const({})
     on_train: Callable[[int], _GeneratorContextManager] = const(nullcontext())
@@ -620,10 +619,8 @@ def train(
     kwargs.update(extra)
 
     vae.to(accelerator.device, dtype=dtype)
-
-    for model in (unet, text_encoder, vae):
-        model.requires_grad_(False)
-        model.eval()
+    vae.requires_grad_(False)
+    vae.eval()
 
     callbacks = strategy.callbacks(
         accelerator=accelerator,
@@ -636,8 +633,6 @@ def train(
         **kwargs,
     )
 
-    callbacks.on_prepare()
-
     loss_step_ = partial(
         loss_step,
         vae,