(ctx, grad_outputs)
| 19 | |
| 20 | @staticmethod |
| 21 | def backward(ctx, grad_outputs): |
| 22 | inputs, targets = ctx.saved_tensors |
| 23 | grad_inputs = None |
| 24 | if ctx.needs_input_grad[0]: |
| 25 | grad_inputs = grad_outputs.mm(ctx.features) |
| 26 | |
| 27 | # momentum update |
| 28 | for x, y in zip(inputs, targets): |
| 29 | ctx.features[y] = ctx.momentum * ctx.features[y] + (1. - ctx.momentum) * x |
| 30 | ctx.features[y] /= ctx.features[y].norm() |
| 31 | |
| 32 | return grad_inputs, None, None, None |
| 33 | |
| 34 | |
| 35 | def cm(inputs, indexes, features, momentum=0.5, device=None): |