I am multiplying a mask to my weight matrix during the forward prop to sparsify my network. However, during backprop, I want the gradient updates to be zero for the masked weights (mask[i] = 0).
Is there any way to do that using register_hook?
Here’s a toy code I have used, but am getting an error.
mask = torch.tensor([0,1,0,1,0])

a.requires_grad = True

b = 3*a

b.retain_grad()

b.register_hook(mask)  

b.mean().backward() 

print(a.grad, b.grad)

Error:
--->b.mean().backward()
    TypeError: 'Tensor' object is not callable

I have tried:
b.register_hook(lambda grad:grad*mask)

which works. Is there any other better way?