| 303 | return out |
| 304 | |
| 305 | def get_loss_iter(self, data, noises=None): |
| 306 | B, D, N = data.shape |
| 307 | t = torch.randint(0, self.diffusion.num_timesteps, size=(B,), device=data.device) |
| 308 | |
| 309 | if noises is not None: |
| 310 | noises[t!=0] = torch.randn((t!=0).sum(), *noises.shape[1:]).to(noises) |
| 311 | |
| 312 | losses = self.diffusion.p_losses( |
| 313 | denoise_fn=self._denoise, data_start=data, t=t, noise=noises) |
| 314 | assert losses.shape == t.shape == torch.Size([B]) |
| 315 | return losses |
| 316 | |
| 317 | def gen_samples(self, shape, device, noise_fn=torch.randn, constrain_fn=lambda x, t:x, |
| 318 | clip_denoised=False, max_timestep=None, |