Upload BertForPrefixMarking.py

BertForPrefixMarking.py

@@ -123,6 +123,7 @@ class BertForPrefixMarking(BertPreTrainedModel):
     def predict(self, sentences: List[str], tokenizer: BertTokenizerFast, padding='longest'):
         # step 1: encode the sentences through using the tokenizer, and get the input tensors + prefix id tensors
         inputs = encode_sentences_for_bert_for_prefix_marking(tokenizer, sentences, padding)
+        inputs = {k:v.to(self.device) for k,v in inputs.items()}
 
         # run through bert
         logits = self.forward(**inputs, return_dict=True).logits