Hi, I have a custom torch.autograd.Function defined like this.
class SourceMemory(Function):
    def __init__(self, M, alpha=0.01):
        super(SourceMemory, self).__init__()
        self.M = M
        self.alpha = alpha

    def forward(self, inputs, targets):
        self.save_for_backward(inputs, targets)
        outputs = inputs.mm(self.M.t())
        return outputs

    def backward(self, grad_outputs):
        inputs, targets = self.saved_tensors
        grad_inputs = None
        if self.needs_input_grad[0]:
            grad_inputs = grad_outputs.mm(self.M)
        for x, y in zip(inputs, targets):
            self.M[y] = self.alpha * self.M[y] + (1. - self.alpha) * x
            self.M[y] = F.normalize(self.M[y], p=2, dim=0)
        return grad_inputs, None


class IntraNet(nn.Module):
    def __init__(self, beta=0.05, alpha=0.01, num_classes=0, num_features=0, weight=None):
        super(IntraNet, self).__init__()
        self.beta = beta
        self.alpha = alpha
        self.weight = weight
        self.register_buffer('M', torch.zeros(num_classes, num_features))

    def forward(self, inputs, targets, features=None, epoch=None):
        self.alpha = self.alpha * epoch
        outputs = SourceMemory(self.M, alpha=self.alpha)(features, targets)
        outputs /= self.beta
        loss = F.cross_entropy(inputs, targets, weight=self.weight)
        return loss, outputs

M is a memory containing the feature. The IntraNet Module is used as a criterion in
  for i, inputs in enumerate(data_loader):
            inputs, pids = self._parse_data(inputs)
            outputs, features = self.model(inputs)
            # this will update the buffer or the memory M
            source_pid_loss, _ = self.inter_criterion(outputs, pids, features=features, epoch=epoch)

Where self.inter_criterion is the IntraNet module. and the model returns a output = [batch_size, num_class] and features = [batch_size, 4096].
The problem is that, backward of SourceMemory is never called. Am I doing something wrong here? I want the buffer M to be filled during the backward pass.
Thank you