Commit
•
3647818
1
Parent(s):
095611c
Upload model
Browse files
modeling_backpack_gpt2_nli.py
CHANGED
|
@@ -56,7 +56,7 @@ class BackpackGPT2NLIModel(GPT2PreTrainedModel):
|
|
| 56 |
|
| 57 |
|
| 58 |
def predict(self, input_ids=None, attention_mask=None):
|
| 59 |
-
logits = self.forward(input_ids, attention_mask, labels=None)
|
| 60 |
p = torch.argmax(logits, axis=1)
|
| 61 |
labels = [self.config.id2label[index] for index in p]
|
| 62 |
return labels
|
|
|
|
| 56 |
|
| 57 |
|
| 58 |
def predict(self, input_ids=None, attention_mask=None):
|
| 59 |
+
logits = self.forward(input_ids, attention_mask, labels=None)['logits']
|
| 60 |
p = torch.argmax(logits, axis=1)
|
| 61 |
labels = [self.config.id2label[index] for index in p]
|
| 62 |
return labels
|