From cda7eba710dfde7b2e67964bcf76cd410c6a4a63 Mon Sep 17 00:00:00 2001
From: Volpeon <git@volpeon.ink>
Date: Mon, 10 Apr 2023 13:42:50 +0200
Subject: Update

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

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

diff --git a/training/strategy/ti.py b/training/strategy/ti.py
index f00045f..08af89d 100644
--- a/training/strategy/ti.py
+++ b/training/strategy/ti.py
@@ -142,7 +142,7 @@ def textual_inversion_strategy_callbacks(
                     checkpoint_output_dir / f"{slugify(token)}_{step}_{postfix}.bin"
                 )
 
-    @on_eval()
+    @torch.no_grad()
     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