Miaoran000 commited on
Commit
dd69ef6
1 Parent(s): 8f6b0b5

support GPU inference

Browse files
Files changed (1) hide show
  1. modeling_hhem_v2.py +1 -1
modeling_hhem_v2.py CHANGED
@@ -59,7 +59,7 @@ class HHEMv2ForSequenceClassification(PreTrainedModel):
59
  tokenizer = self.tokenzier
60
  pair_dict = [{'text1': pair[0], 'text2': pair[1]} for pair in text_pairs]
61
  inputs = tokenizer(
62
- [self.prompt.format(**pair) for pair in pair_dict], return_tensors='pt', padding=True)
63
  self.t5.eval()
64
  with torch.no_grad():
65
  outputs = self.t5(**inputs)
 
59
  tokenizer = self.tokenzier
60
  pair_dict = [{'text1': pair[0], 'text2': pair[1]} for pair in text_pairs]
61
  inputs = tokenizer(
62
+ [self.prompt.format(**pair) for pair in pair_dict], return_tensors='pt', padding=True).to(self.t5.device)
63
  self.t5.eval()
64
  with torch.no_grad():
65
  outputs = self.t5(**inputs)