From e47c97a45129de2b2dc59a52bf76861dff5d3bdb Mon Sep 17 00:00:00 2001
From: Patrick von Platen <patrick.v.platen@gmail.com>
Date: Tue, 28 Jun 2022 23:05:08 +0000
Subject: [PATCH] no inference moed doesn't always work

---
 README.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/README.md b/README.md
index 8f6db61267..7ca8d7b23b 100644
--- a/README.md
+++ b/README.md
@@ -161,7 +161,7 @@ for t in tqdm.tqdm(reversed(range(num_inference_steps)), total=num_inference_ste
     # 1. predict noise residual
     orig_t = len(noise_scheduler) // num_inference_steps * t
 
-    with torch.inference_mode():
+    with torch.no_grad():
         residual = unet(image, orig_t)
 
     # 2. predict previous mean of image x_t-1