PyTorch - backward () не вызывается в пользовательском слое

Question

Что я собираюсь сделать, это изменить вес в Conv2d после loss.backward () и перед optimizer.step ().Одним из решений является изменение весов в соответствующих слоях после loss.backward ().Я просто хочу создать пользовательский слой, чтобы сохранить функцию train () в чистоте.Вот фрагмент кода.

class CustomConv(nn.Conv2d):
    def __init__(self, *kargs, **kwargs):
        super(CustomConv, self).__init__(*kargs, **kwargs)

    def forward(self, input):
        input = WeightModifier(self.weight)(input)
        out = nn.functional.conv2d(
            input,
            self.weight,
            None,
            self.stride,
            self.padding,
            self.dilation,
            self.groups)
        return out

class WeightModifier(torch.autograd.Function):
    def __init__(self, weight):
        self.weight = weight

    def forward(self, input):
        return input.clone()

    def backward(self, grad):
        # Some op to change self.weight.
        # In other words, change weights in the following conv2d
        # after loss.backward() and before optimizer.step()
        return grad.clone()

hidden = CustomConv()(input_tensor)  # backward not called 
hidden = CustomConv()(hidden)
loss = cal_loss(hidden, gt)
optimizer.zero_grad()
loss.backward()
optimizer.step()

Проблема заключается в том, что функция backward () для WeightModifier в первом CustomConv не вызывается (во втором CustomConv вызывается).Я предполагаю, что причина в том, что Pytorch находит, что input_tensor не требует градиента, а слой WeightModifier не имеет каких-либо параметров, поэтому он пропускает этот метод backward ().Является ли какой-либо метод заставить или обмануть Pytorch для выполнения функции backward ()?

Спасибо!