Update: I modified the code a bit and printed the the incoming gradient shape and shape of the mask I am passing in the loop. It seems that the way backward hooks are called, it always takes the first value ie ‘linear1’(=name)  so same mask is getting passed to all the gradients resulting in mis-match matrix size error. Here’s the output I got:
shape of grad torch.Size([10, 100])  mask-shape torch.Size([10, 100])

shape of grad torch.Size([100, 300])  mask-shape torch.Size([10, 100])

shape of grad torch.Size([300, 784])  mask-shape torch.Size([10, 100])

shape of grad torch.Size([10, 100])  mask-shape torch.Size([10, 100])

shape of grad torch.Size([100, 300])  mask-shape torch.Size([10, 100])

shape of grad torch.Size([300, 784])  mask-shape torch.Size([10, 100])

and this repeats on…