Sorry, I forgot to update, I was able to finally make the backward hook work but manually. I am trying to automate the process, for that I have a mask_dict which has key, value pairs as layer_name, bool_mask. I am trying to loop through the layer modules and hook using mask_dict keys but it doesn’t work. There’s no particular reason for using bool values, I just chose them as mask is boolean in nature.
def random_mask(model):
    mask_dict = {}
    for name, module in model.named_modules():
        if 'linear' in name:
            size = tuple(module.weight.shape)
            mask_dict[name] = (torch.randint(0, 2, size).bool().to(device='cuda'))
    return mask_dict        
               

def apply_mask(model, mask_dict):
    for name, module in model.named_modules():
        if 'linear' in name:
            module.weight.data *= mask_dict[name]
            #checking if names of layers and it's tensor shapes match with that of masks
            print('module name is:', name, 'and weight size is:', module.weight.size()) 
            print('corresponding tensor is:', mask_dict[name].shape) #matching shapes for multiplication

            module.weight.register_hook(lambda x: x*mask_dict[name])  #<--loop update doesn't work
           
#### manual update
#     model.linear1.weight.register_hook(lambda x: x*mask_dict['linear1'])
#     model.linear2.weight.register_hook(lambda x: x*mask_dict['linear2'])
#     model.linear3.weight.register_hook(lambda x: x*mask_dict['linear3'])
   

mask = random_mask(model)
#print(mask)
apply_mask(model, mask)


module name is: linear1 and weight size is: torch.Size([300, 784])
corresponding tensor is: torch.Size([300, 784]) 

module name is: linear2 and weight size is: torch.Size([100, 300])
corresponding tensor is: torch.Size([100, 300])

module name is: linear3 and weight size is: torch.Size([10, 100])
corresponding tensor is: torch.Size([10, 100])

It works perfect when I un-comment the manual update line and remove the loop update but not in it’s current state. The error I am getting is of tensor size mismatch for performing element-wise multiplication.