Thanks for your help. It is so useful to get familiar with hook in pytorch. I only treated it as a intermediate output reader in the past.
I consider the complete version will be:
def modify_input(idx, mask):
    def hook(module, input):
        if not module.training:
            input[idx] = mask
    return hook
hook = net[1].register_forward_pre_hook(modify_input(idx, mask))