ErfanMoosaviMonazzah commited on
Commit
eb0bbcd
1 Parent(s): aa65e77

Upload model

Browse files
Files changed (1) hide show
  1. modeling_backpack_gpt2_nli.py +1 -1
modeling_backpack_gpt2_nli.py CHANGED
@@ -59,5 +59,5 @@ class BackpackGPT2NLIModel(GPT2PreTrainedModel):
59
  logits = self.forward(input_ids, attention_mask, labels=None)['logits']
60
  p = torch.argmax(logits, axis=1)
61
  labels = [self.config.id2label[index.item()] for index in p]
62
- return labels
63
 
 
59
  logits = self.forward(input_ids, attention_mask, labels=None)['logits']
60
  p = torch.argmax(logits, axis=1)
61
  labels = [self.config.id2label[index.item()] for index in p]
62
+ return {'labels':labels, 'logits':logits}
63