Thanks a lot for your response. I have another question. Lets say I am introducing a parameter c like this:
a = torch.ones(5)
a.requires_grad = True
mask = torch.tensor([1,0,0,1,0])
b = 3*a
c = 2*b

b.retain_grad()
c.retain_grad()

b.register_hook(lambda grad: grad*mask)  

c.mean().backward() 
print(a.grad,b.grad,c.grad)

Here, only a.grad gets masked. Both a.grad and b.grad get masked only when I use
c.register_hook(lambda grad: grad*mask)  

Why is that? Shouldnt b.grad get masked even in the first scenario?