Visualize feature map

You could either return this activation in your forward (or store in a list or dict), if it’s a functional call.
On the other hand, if this Mul operation is performed in an nn.Module, you could use forward hooks directly.