1 files changed, 1 insertions, 1 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}...")