1
0
mirror of https://github.com/huggingface/diffusers.git synced 2026-01-27 17:22:53 +03:00

use inference_mode

This commit is contained in:
Kashif Rasul
2022-06-16 17:54:47 +02:00
parent d2940c23fe
commit cf3fdb8479

View File

@@ -147,9 +147,9 @@ eta = 0.0 # <- deterministic sampling
for t in tqdm.tqdm(reversed(range(num_inference_steps)), total=num_inference_steps):
# 1. predict noise residual
orig_t = noise_scheduler.get_orig_t(t, num_inference_steps)
with torch.no_grad():
residual = unet(image, orig_t)
orig_t = noise_scheduler.get_orig_t(t, num_inference_steps)
with torch.inference_mode():
residual = unet(image, orig_t)
# 2. predict previous mean of image x_t-1
pred_prev_image = noise_scheduler.step(residual, image, t, num_inference_steps, eta)