Skip to content

GitLab

Commit cf3fdb84 authored by Kashif Rasul's avatar Kashif Rasul
Browse files

use inference_mode

parent d2940c23
Hide whitespace changes
Inline Side-by-side
Showing with 3 additions and 3 deletions
README.md
View file @ cf3fdb84
...@@ -147,9 +147,9 @@ eta = 0.0 # <- deterministic sampling ...@@ -147,9 +147,9 @@ eta = 0.0 # <- deterministic sampling
for t in tqdm.tqdm(reversed(range(num_inference_steps)), total=num_inference_steps): for t in tqdm.tqdm(reversed(range(num_inference_steps)), total=num_inference_steps):
# 1. predict noise residual # 1. predict noise residual
orig_t = noise_scheduler.get_orig_t(t, num_inference_steps) orig_t = noise_scheduler.get_orig_t(t, num_inference_steps)
with torch.no_grad(): with torch.inference_mode():
residual = unet(image, orig_t) residual = unet(image, orig_t)
# 2. predict previous mean of image x_t-1 # 2. predict previous mean of image x_t-1
pred_prev_image = noise_scheduler.step(residual, image, t, num_inference_steps, eta) pred_prev_image = noise_scheduler.step(residual, image, t, num_inference_steps, eta)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment