Fix embeddings not working with --gpu-only

im-stuff · Jun 30, 2023 · 9920367 · 9920367
1 parent 62db116
commit 9920367
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/comfy/sd1_clip.py b/comfy/sd1_clip.py
@@ -95,7 +95,7 @@ def set_up_textual_embeddings(self, tokens, current_embeds):
             out_tokens += [tokens_temp]
 
         if len(embedding_weights) > 0:
-            new_embedding = torch.nn.Embedding(next_new_token, current_embeds.weight.shape[1])
+            new_embedding = torch.nn.Embedding(next_new_token, current_embeds.weight.shape[1], device=self.device)
             new_embedding.weight[:token_dict_size] = current_embeds.weight[:]
             n = token_dict_size
             for x in embedding_weights: