fix self conditioning

lucidrains · Jan 22, 2024 · 32310c3 · ChoCho66 · Jan 23, 2024 · 32310c3
1 parent 242e4b6
commit 32310c3
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py b/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py
@@ -785,7 +785,7 @@ def p_losses(self, x_start, t, noise = None, offset_noise_strength = None):
 
         x_self_cond = None
         if self.self_condition and random() < 0.5:
-            with torch.inference_mode():
+            with torch.no_grad():
                 x_self_cond = self.model_predictions(x, t).pred_x_start
                 x_self_cond.detach_()
 

diff --git a/denoising_diffusion_pytorch/version.py b/denoising_diffusion_pytorch/version.py
@@ -1 +1 @@
-__version__ = '1.9.5'
+__version__ = '1.9.6'