I think this would be a job for forward hooks. This 4 thread might be of interest to you. Here 4 are the docs on how to use forward hooks.
In essence how forward hooks work:

Create a hooking function which has inputs function(module, input, output)
Register the forward hook to some module (Note: all pytorch fuctions and layers subclass the Module class)
When the forward function is called on the module which has a forward hook registered the function you registered will be caled with the parameters: (module, module_input, module_output)
Within the function you can intercept the module_otuput and do whatever visualizations you like on the model_output values. Note: The hooking functions should not modify the output, so I suggest you copy the data to some other tensor if you want to process the data in some way before visualizing