Sure
def get_hook(param_idx):
   def hook(grad):
        grad = grad.clone() # NEVER change the given grad inplace
        # Assumes 1D but can be generalized
        for i in grad.size(0):
            if should_be_disabled[param_idx][i]:
                grad[i] = 0
        return grad
    return hook

for j, p in enumerate(model.parameters()):
    p.register_hook(get_hook(j))