You could use the self.training flag in your custom module to switch between the different behaviors:

apply your mask and scale the activations if self.training

disable the mask and just forward the activations if self.training == False


Let me know, if that works for you.