Update models/unet.py

models/unet.py

@@ -864,6 +864,10 @@ class MotionCLR(nn.Module):
         self.unet = self.unet.cuda()
 
     def encode_text(self, raw_text, device):
+        self.clip_model.token_embedding = self.clip_model.token_embedding.to(device)
+        self.clip_model.positional_embedding = self.clip_model.positional_embedding.to(device)
+        self.clip_model.transformer = self.clip_model.transformer.to(device)
+        self.clip_model.ln_final = self.clip_model.ln_final.to(device)
         with torch.no_grad():
             texts = clip.tokenize(raw_text, truncate=True).to(
                 device