shreyd:

Ideally, the operation p*c would happen before p is used in the forward pass. Then, during the backwards pass, both p and c would receive a gradient wrt the loss.


Hello,
I tried to define hooks to achieve it but encountered some problem.
c = nn.Parameter(torch.randn(1), requires_grad=True)
opt = torch.optim.SGD(model.parameters(), lr=0.001)
opt.add_param_group({'params': c})

def hooks(module, input):
    module.weight = nn.Parameter(module.weight * c)

for module in model.children():
    if isinstance(module, nn.Conv2d):
        module.register_forward_pre_hook(hooks)

c is defined as nn.Parameter and if we want to modify module.weight we must pass nn.Parameter or None to it, then registered it to each child_module.
But, c is still not a learnable param, and what I found is:

module.weight is leaf variable, it is okay since it used to update.
c is also leaf variable, it means that gradient flow will not pass back to c? and c is not included in the computation graph?