| 109 | |
| 110 | |
| 111 | class WithinEmbedContrast(torch.nn.Module): |
| 112 | def __init__(self, loss: Loss, **kwargs): |
| 113 | super(WithinEmbedContrast, self).__init__() |
| 114 | self.loss = loss |
| 115 | self.kwargs = kwargs |
| 116 | |
| 117 | def forward(self, h1, h2): |
| 118 | l1 = self.loss(anchor=h1, sample=h2, **self.kwargs) |
| 119 | l2 = self.loss(anchor=h2, sample=h1, **self.kwargs) |
| 120 | return (l1 + l2) * 0.5 |