(self, x, *args, **kwargs)
| 418 | return loss, loss_dict |
| 419 | |
| 420 | def forward(self, x, *args, **kwargs): |
| 421 | # b, c, h, w, device, img_size, = *x.shape, x.device, self.image_size |
| 422 | # assert h == img_size and w == img_size, f'height and width of image must be {img_size}' |
| 423 | t = torch.randint(0, self.num_timesteps, (x.shape[0],), device=self.device).long() |
| 424 | return self.p_losses(x, t, *args, **kwargs) |
| 425 | |
| 426 | def get_input(self, batch, k): |
| 427 | x = batch[k] |