2 files changed, 2 insertions, 2 deletions

diff --git a/trainer/base.py b/trainer/base.py
index e700dd6..1f85e71 100644
--- a/trainer/base.py
+++ b/trainer/base.py
@@ -74,7 +74,7 @@ class Checkpointer():
     def checkpoint(self, step: int, postfix: str):
         pass
 
-    @torch.inference_mode()
+    @torch.no_grad()
     def save_samples(self, step: int):
         print(f"Saving samples for step {step}...")
 
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()