| 98 | self.register_buffer('features', torch.zeros(num_samples, num_features)) |
| 99 | |
| 100 | def forward(self, inputs, targets): |
| 101 | if self.device is not None: |
| 102 | inputs = F.normalize(inputs, dim=1).to(self.device) |
| 103 | targets = targets.to(self.device) |
| 104 | else: |
| 105 | inputs = F.normalize(inputs, dim=1).cuda() |
| 106 | targets = targets.cuda() |
| 107 | if self.use_hard: |
| 108 | outputs = cm_hard(inputs, targets, self.features, self.momentum, self.device) |
| 109 | else: |
| 110 | outputs = cm(inputs, targets, self.features, self.momentum, self.device) |
| 111 | |
| 112 | outputs /= self.temp |
| 113 | loss = F.cross_entropy(outputs, targets) |
| 114 | return loss |