Update models/unet.py

models/unet.py

@@ -856,14 +856,13 @@ class MotionCLR(nn.Module):
             log_attn=self.vis_attn,
             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)
 
     def encode_text(self, raw_text, device):
-        print("00000000")
-        print(device)
-        print(next(self.clip_model.parameters()).device)
-        print("00000000")
-        self.clip_model = self.clip_model.to(device)
-        print("00000000")
         with torch.no_grad():
             texts = clip.tokenize(raw_text, truncate=True).to(
                 device