From 21d70916f66e74a87c631a06b70774954b085b48 Mon Sep 17 00:00:00 2001 From: Volpeon Date: Fri, 7 Apr 2023 14:14:00 +0200 Subject: Fix --- training/strategy/lora.py | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) (limited to 'training/strategy/lora.py') diff --git a/training/strategy/lora.py b/training/strategy/lora.py index 6730dc9..80ffa9c 100644 --- a/training/strategy/lora.py +++ b/training/strategy/lora.py @@ -64,17 +64,17 @@ def lora_strategy_callbacks( image_size=sample_image_size, ) - def on_accum_model(): - return unet - @contextmanager def on_train(epoch: int): - tokenizer.train() + unet.train() text_encoder.train() + tokenizer.train() yield @contextmanager def on_eval(): + unet.eval() + text_encoder.eval() tokenizer.eval() yield @@ -152,7 +152,6 @@ def lora_strategy_callbacks( torch.cuda.empty_cache() return TrainingCallbacks( - on_accum_model=on_accum_model, on_train=on_train, on_eval=on_eval, on_before_optimize=on_before_optimize, -- cgit v1.2.3-54-g00ecf