From 8c068963d4b67c6b894e720288e5863dade8d6e6 Mon Sep 17 00:00:00 2001
From: Volpeon <git@volpeon.ink>
Date: Sat, 31 Dec 2022 13:09:04 +0100
Subject: Fixes

---
 models/clip/embeddings.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

(limited to 'models/clip/embeddings.py')

diff --git a/models/clip/embeddings.py b/models/clip/embeddings.py
index 7d63ffb..f82873e 100644
--- a/models/clip/embeddings.py
+++ b/models/clip/embeddings.py
@@ -74,7 +74,7 @@ class ManagedCLIPTextEmbeddings(CLIPTextEmbeddings):
 
     def get_embed(self, input_ids: Union[list[int], torch.LongTensor]):
         if isinstance(input_ids, list):
-            input_ids = torch.tensor(input_ids)
+            input_ids = torch.tensor(input_ids, device=self.token_embedding.weight.device)
 
         mask = torch.isin(input_ids, torch.tensor(self.temp_token_ids, device=input_ids.device))
 
-- 
cgit v1.2.3-70-g09d2