(self, h, g, batch=None, hn=None, extra_pos_mask=None, extra_neg_mask=None)
| 24 | self.kwargs = kwargs |
| 25 | |
| 26 | def forward(self, h, g, batch=None, hn=None, extra_pos_mask=None, extra_neg_mask=None): |
| 27 | if batch is None: # for single-graph datasets |
| 28 | assert hn is not None |
| 29 | anchor, sample, pos_mask, neg_mask = self.sampler(anchor=g, sample=h, neg_sample=hn) |
| 30 | else: # for multi-graph datasets |
| 31 | assert batch is not None |
| 32 | anchor, sample, pos_mask, neg_mask = self.sampler(anchor=g, sample=h, batch=batch) |
| 33 | |
| 34 | pos_mask, neg_mask = add_extra_mask(pos_mask, neg_mask, extra_pos_mask, extra_neg_mask) |
| 35 | loss = self.loss(anchor=anchor, sample=sample, pos_mask=pos_mask, neg_mask=neg_mask, **self.kwargs) |
| 36 | return loss |
| 37 | |
| 38 | |
| 39 | class DualBranchContrast(torch.nn.Module): |
nothing calls this directly
no test coverage detected