From fc11c86142915d6c3935d28a3321b3ae91b613ef Mon Sep 17 00:00:00 2001
From: Volpeon <git@volpeon.ink>
Date: Sat, 14 Jan 2023 22:03:01 +0100
Subject: Update

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

(limited to 'trainer/ti.py')

diff --git a/trainer/ti.py b/trainer/ti.py
index 15cf747..388acd3 100644
--- a/trainer/ti.py
+++ b/trainer/ti.py
@@ -42,7 +42,7 @@ class TextualInversionCheckpointer(Checkpointer):
                     checkpoints_path.joinpath(f"{slugify(token)}_{step}_{postfix}.bin")
                 )
 
-    @torch.inference_mode()
+    @torch.no_grad()
     def save_samples(self, step):
         ema_context = self.ema_embeddings.apply_temporary(
             self.text_encoder.text_model.embeddings.temp_token_embedding.parameters()
-- 
cgit v1.2.3-54-g00ecf