I think the solution is
def freeze_weights(model, hooks = [], gradient_mask_weights ={}):
    ...
    ...
    h = module.weight.register_hook(lambda grad, n=name: grad.mul_(gradient_mask_dict[n]))
    ...
    ...
    return hooks, gradient_mask_weights