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/base.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

(limited to 'trainer/base.py')

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}...")
 
-- 
cgit v1.2.3-54-g00ecf