What I am going to do is modifying weight in Conv2d after loss.backward() and before optimizer.step(). One solution is to modify weights in corresponding layers after loss.backward(). I just want to make a custom layer to do this keeping train() function clean. Here is the code snippet:
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()

The problem is the backward() of WeightModifier in the first CustomConv is not called (that in second CustomConv is called). I guess the reason is that Pytorch finds that input_tensor does not require gradient and layer WeightModifier does not have any parameters. Is the any method to force or “cheat” Pytorch to execute the backward()?
Thanks.