Update models/unet.py

models/unet.py

@@ -857,10 +857,10 @@ class MotionCLR(nn.Module):
             edit_config=edit_config,
         )
         
-        self.embed_text = self.embed_text.to(device)
-        self.textTransEncoder = self.textTransEncoder.to(device)
-        self.text_ln = self.text_ln.to(device)
-        self.unet = self.unet.to(device)
+        self.embed_text = self.embed_text.cuda()
+        self.textTransEncoder = self.textTransEncoder.cuda()
+        self.text_ln = self.text_ln.cuda()
+        self.unet = self.unet.cuda()
 
     def encode_text(self, raw_text, device):
         with torch.no_grad():