I am trying to implement the LotteryTicket Hypothesis 2 paper in Pytorch. For freezing the individual weights, I have created a mask_dict with layer name as key as boolean mask as the value. To freeze, I wanted to use backward hooks via register_hook and multiply incoming gradients by the mask to set some values to zero. The problem is whatever I try, I always a RumtimeError saying the size of tensors in not compatible.
Here’s the code for creating mask
for name, t in learn.model.named_parameters():
    if 'weight' in name:
        bottomk = torch.topk(torch.abs(t).view(-1), k=round(pr_value*t.nelement()), largest=False)
        clone = t.clone().detach()
        clone.view(-1)[bottomk.indices] = 0
        mask[name] = clone.bool()

I thought that maybe since incoming gradients come in reverse, it results in a mismatch but even switching the mask values doesn’t help. Below is my attempt,
for name, module in learn.model.named_modules():
    if 'linear' in name:
        key = name+'.weight'       # since layer keys have names layer.weight
        module.weight.register_hook(lambda x: x*mask[key])