Hi albanD,
Thanks. I still don’t make it work. As you said, I tried something like this:
class WeightModifier(torch.autograd.Function):
    @staticmethod
    def forward(ctx, input, weight):
        ctx.save_for_backward(weight)
        return input.clone()

    @staticmethod
    def backward(ctx, grad_output):
        weight, = ctx.saved_tensors
        weight.data.copy_(weight.org)
        return grad_output.clone()

which is called by input = WeightModifier.apply(input, self.weight). However Pytorch requires backward() to return gradients for both input and weight, whereas I want it to return gradient for input only.