That won’t be optimized away. If you really want to optimize it this should work:
nonlinearities = [F.relu, F.tanh, F.sigmoid]
masks = generate_masks(input) # a list of masks
preactivation = module(input) 
# preallocate output
# IMPORTANT: it shouldn't require grad!
# you don't care about the grad w.r.t the original content!
output = Variable(preactivation.new(preactivation.size()))
for nonlinearity, mask in zip(nonlinearities, masks):
    output[mask] = nonlinearity(preactivation[mask])