Lazyhope commited on
Commit
8fc8456
1 Parent(s): c0f0c0a

Add source id to self.device

Browse files
Files changed (1) hide show
  1. 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)