Update models/unet.py

models/unet.py

@@ -808,7 +808,7 @@ class MotionCLR(nn.Module):
         # text encoder
         self.embed_text = nn.Linear(clip_dim, text_latent_dim)
         self.clip_version = clip_version
-        self.clip_model = self.load_and_freeze_clip(clip_version)
+        self.clip_model = self.load_and_freeze_clip(clip_version).cuda()
         textTransEncoderLayer = nn.TransformerEncoderLayer(
             d_model=text_latent_dim,
             nhead=text_num_heads,