Update models/unet.py

models/unet.py

@@ -840,16 +840,14 @@ class MotionCLR(nn.Module):
     def encode_text(self, raw_text, device):
         print("00000000")
         print(device)
-        print(self.clip_model.device)
+        print(next(self.clip_model.parameters()).device)
         print("00000000")
         with torch.no_grad():
             texts = clip.tokenize(raw_text, truncate=True).to(
                 device
             )  # [bs, context_length] # if n_tokens > 77 -> will truncate
-            x = self.clip_model.token_embedding(texts).type(
-                self.clip_model.dtype
-            )  # [batch_size, n_ctx, d_model]
-            x = x + self.clip_model.positional_embedding.type(self.clip_model.dtype)
+            x = self.clip_model.token_embedding(texts).type(self.clip_model.dtype).to(device)  # [batch_size, n_ctx, d_model]
+            x = x + self.clip_model.positional_embedding.type(self.clip_model.dtype).to(device)
             x = x.permute(1, 0, 2)  # NLD -> LND
             x = self.clip_model.transformer(x)
             x = self.clip_model.ln_final(x).type(