Add source id to self.device
Browse files- pipeline.py +1 -1
pipeline.py
CHANGED
@@ -144,7 +144,7 @@ class RepoEmbeddingPipeline(Pipeline):
|
|
144 |
+ [tokenizer.sep_token]
|
145 |
)
|
146 |
tokens_id = tokenizer.convert_tokens_to_ids(tokens)
|
147 |
-
source_ids = torch.tensor([tokens_id])
|
148 |
|
149 |
token_embeddings = self.model(source_ids)[0]
|
150 |
sentence_embeddings = token_embeddings.mean(dim=1)
|
|
|
144 |
+ [tokenizer.sep_token]
|
145 |
)
|
146 |
tokens_id = tokenizer.convert_tokens_to_ids(tokens)
|
147 |
+
source_ids = torch.tensor([tokens_id]).to(self.device)
|
148 |
|
149 |
token_embeddings = self.model(source_ids)[0]
|
150 |
sentence_embeddings = token_embeddings.mean(dim=1)
|