Feature Request: pytorch and tensorflow

[jjennings955]

Nice library, I've rewritten this same function so many times, it would be great if it was standardized for all the common libs (maybe even submit a pr to matplotlib?)

Specifically it would be nice if this library transparently handed torch and tensorflow tensors and different dimension layouts (NHWC) vs (NCHW)

To use matplotlib with pytorch you always have to do something like tensor.cpu().detach().numpy().permute(0, 2, 3, 1)
before any plotting.

It makes code very ugly.

Also, in my experience using subplots is super slow and it's much better to lay things out yourself in a numpy array (haven't benchmarked your version though).

pytorch has its own version of this but it would be nice to have a common clean API I could use everywhere.

I can maybe submit a pr for this if you're interested.

[jjennings955]

This is a good implementation without using subplots
https://github.com/keras-team/keras/blob/7a39b6c62d43c25472b2c2476bd2a8983ae4f682/examples/conv_filter_visualization.py#L178