You could register a hook directly on the parameter, e.g. as:
model = nn.Conv2d(3, 6, 3, 1, 1)
mask = torch.randint(0, 2, (6, 3, 3, 3)).float()
model.weight.register_hook(lambda x: x * mask)

model(torch.randn(1, 3, 4, 4)).mean().backward()
print(model.weight.grad)