I see… so I would do something like this
out_ = model.layer1.register_hook(lambda grad: torch.t(torch.mm(M, torch.t(grad))))
out_2 = model.layer2.register_hook(lambda grad: torch.t(torch.mm(M, torch.t(grad))))
Is it right?
but then how would I make the optimizer to use those updated weights?
e.g SGD optimizer