Solved by spanev in post #2 
              
              
                Hi @zeke_wang 
You should be able to do it with forward pre hooks. 
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 loo…