From e9dc712268e45d30451fc6fee8626a0a8af7ccdc Mon Sep 17 00:00:00 2001
From: Volpeon <git@volpeon.ink>
Date: Mon, 10 Apr 2023 12:57:21 +0200
Subject: Fix sample gen: models sometimes weren't in eval mode

---
 training/strategy/lora.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

(limited to 'training/strategy/lora.py')

diff --git a/training/strategy/lora.py b/training/strategy/lora.py
index 73ec8f2..0f72a17 100644
--- a/training/strategy/lora.py
+++ b/training/strategy/lora.py
@@ -146,7 +146,7 @@ def lora_strategy_callbacks(
         if torch.cuda.is_available():
             torch.cuda.empty_cache()
 
-    @torch.no_grad()
+    @on_eval()
     def on_sample(step):
         unet_ = accelerator.unwrap_model(unet, keep_fp32_wrapper=True)
         text_encoder_ = accelerator.unwrap_model(text_encoder, keep_fp32_wrapper=True)
-- 
cgit v1.2.3-54-g00ecf