I think the better way would be to use forward hooks as explained in this post. Your current approach could break, if you are using functional API calls inside the forward and you would need to add them.
Once you have the tensors, you could visualize them via e.g. matplotlib.