Update modeling_bertchunker.py

modeling_bertchunker.py

@@ -41,7 +41,7 @@ class BertChunker(PreTrainedModel):
     # slide context window
         MAX_TOKENS=255
         tokens=tokenizer(text, return_tensors="pt",truncation=False)
-        input_ids=tokens['input_ids']
+        input_ids=tokens['input_ids'].to(self.device)
         attention_mask=tokens['attention_mask'][:,0:MAX_TOKENS]
         attention_mask=attention_mask.to(self.device)
         CLS=input_ids[:,0].unsqueeze(0)
@@ -60,7 +60,7 @@ class BertChunker(PreTrainedModel):
 
             ids=ids.to(self.device)
             
-            output=self(input_ids=ids,attention_mask=torch.ones(1, ids.shape[1]))
+            output=self(input_ids=ids,attention_mask=torch.ones(1, ids.shape[1],device=self.device))
             logits = output['logits'][:, 1:-1,:]
 
             chunk_probabilities = F.softmax(logits, dim=-1)[:,:,1]