add device

diffloss.py

@@ -35,12 +35,12 @@ class DiffLoss(nn.Module):
     def sample(self, z, temperature=1.0, cfg=1.0):
         # diffusion loss sampling
         if not cfg == 1.0:
-            noise = torch.randn(z.shape[0] // 2, self.in_channels)
+            noise = torch.randn(z.shape[0] // 2, self.in_channels).to(device)
             noise = torch.cat([noise, noise], dim=0)
             model_kwargs = dict(c=z, cfg_scale=cfg)
             sample_fn = self.net.forward_with_cfg
         else:
-            noise = torch.randn(z.shape[0], self.in_channels)
+            noise = torch.randn(z.shape[0], self.in_channels).to(device)
             model_kwargs = dict(c=z)
             sample_fn = self.net.forward