Hi @zeke_wang
You should be able to do it with forward pre hooks 15.
Let’s say you have the following model:
net = nn.Sequential(nn.Linear(16,4),
                    nn.Linear(4, 16))

If you want to modify the activation maps that goes to the second Linear, and only in eval mode, your hook will look like this:
def hook_fn(module, input)
    if not module.training:
        for i in input:
            i[mask] = 0

and to register it:
hook = net[1].register_forward_pre_hook(hook_fn)