From 0316063e2d873b52994d582ffb1c529e75dd6aa7 Mon Sep 17 00:00:00 2001 From: James Betker Date: Fri, 10 Jun 2022 15:37:02 -0600 Subject: [PATCH] . --- codes/trainer/injectors/gaussian_diffusion_injector.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/codes/trainer/injectors/gaussian_diffusion_injector.py b/codes/trainer/injectors/gaussian_diffusion_injector.py index 7ddceeb3..9b2e794d 100644 --- a/codes/trainer/injectors/gaussian_diffusion_injector.py +++ b/codes/trainer/injectors/gaussian_diffusion_injector.py @@ -4,7 +4,7 @@ import torch from torch.cuda.amp import autocast from models.diffusion.gaussian_diffusion import get_named_beta_schedule -from models.diffusion.resample import create_named_schedule_sampler, LossAwareSampler, DeterministicSampler +from models.diffusion.resample import create_named_schedule_sampler, LossAwareSampler, DeterministicSampler, LossSecondMomentResampler from models.diffusion.respace import space_timesteps, SpacedDiffusion from trainer.inject import Injector from utils.util import opt_get