Fix clip vision lowvram mode not working.

kamilsulima · Dec 27, 2023 · c782144 · c782144
1 parent e478b17
commit c782144
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/comfy/clip_model.py b/comfy/clip_model.py
@@ -151,7 +151,7 @@ def __init__(self, embed_dim, num_channels=3, patch_size=14, image_size=224, dty
 
     def forward(self, pixel_values):
         embeds = self.patch_embedding(pixel_values).flatten(2).transpose(1, 2)
-        return torch.cat([self.class_embedding.expand(pixel_values.shape[0], 1, -1), embeds], dim=1) + self.position_embedding.weight
+        return torch.cat([self.class_embedding.to(embeds.device).expand(pixel_values.shape[0], 1, -1), embeds], dim=1) + self.position_embedding.weight.to(embeds.device)
 
 
 class CLIPVision(torch.nn.Module):