(self, batch, batch_idx)
| 79 | return {'loss': loss} |
| 80 | |
| 81 | def test_step(self, batch, batch_idx): |
| 82 | pos_ids, pos_mask, neg_ids, neg_mask = batch |
| 83 | |
| 84 | neg_ids = neg_ids.view(-1, neg_ids.shape[-1]) |
| 85 | neg_mask = neg_mask.view(-1, neg_mask.shape[-1]) |
| 86 | |
| 87 | pos_outputs = self(pos_ids, attention_mask=pos_mask, labels=torch.ones(pos_ids.shape[0], dtype=torch.uint8).to(pos_ids.get_device())) |
| 88 | neg_outputs = self(neg_ids, attention_mask=neg_mask, labels=torch.zeros(neg_ids.shape[0], dtype=torch.uint8).to(neg_ids.get_device())) |
| 89 | |
| 90 | loss_scale = 1.0 |
| 91 | loss = pos_outputs.loss + loss_scale * neg_outputs.loss |
| 92 | |
| 93 | pos_logits = pos_outputs.logits |
| 94 | pos_preds = torch.argmax(pos_logits, axis=1) |
| 95 | self.test_acc(pos_preds.cpu(), torch.ones(pos_ids.shape[0], dtype=torch.uint8).cpu()) |
| 96 | |
| 97 | neg_logits = neg_outputs.logits |
| 98 | neg_preds = torch.argmax(neg_logits, axis=1) |
| 99 | self.test_acc(neg_preds.cpu(), torch.zeros(neg_ids.shape[0], dtype=torch.uint8).cpu()) |
| 100 | |
| 101 | |
| 102 | self.log('test_acc', self.test_acc) |
nothing calls this directly
no outgoing calls
no test coverage detected