From d85b8bbe2e00242575e38c5df46014ef0e860ed4 Mon Sep 17 00:00:00 2001 From: Kashif Rasul Date: Fri, 16 Sep 2022 17:19:00 +0200 Subject: [PATCH] fix torch.cumprod --- denoising_diffusion_pytorch/denoising_diffusion_pytorch.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py b/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py index c8d0086..6e98474 100644 --- a/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py +++ b/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py @@ -450,7 +450,7 @@ class GaussianDiffusion(nn.Module): raise ValueError(f'unknown beta schedule {beta_schedule}') alphas = 1. - betas - alphas_cumprod = torch.cumprod(alphas, axis=0) + alphas_cumprod = torch.cumprod(alphas, dim=0) alphas_cumprod_prev = F.pad(alphas_cumprod[:-1], (1, 0), value = 1.) timesteps, = betas.shape