use inference_mode
This commit is contained in:
parent
d2940c23fe
commit
cf3fdb8479
|
@ -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)
|
||||
|
|
Loading…
Reference in New Issue