fix device issue

unet.py

@@ -411,6 +411,10 @@ class ConditionalUnet(nn.Module):
 
         self.class_embedding = nn.Embedding(num_classes + 1, unet.starting_channels, padding_idx=0)
     
+    def to(self, device):
+        self.device = device
+        return super().to(device)
+    
     def forward(self, x, t, cond=None):
         cond = torch.tensor([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0])
         cond = cond.unsqueeze(0)