Update models/unet.py

models/unet.py

@@ -879,8 +879,11 @@ class MotionCLR(nn.Module):
                 self.clip_model.dtype
             )  # [len, batch_size, 512]
 
+        self.embed_text = self.embed_text.to(device)
         x = self.embed_text(x)  # [len, batch_size, 256]
+        self.textTransEncoder = self.textTransEncoder.to(device)
         x = self.textTransEncoder(x)
+        self.text_ln = self.text_ln.to(device)
         x = self.text_ln(x)
         
         # T, B, D -> B, T, D