Simplify matplotlib plotting

[iancze]

I discovered by accident that a recent version of PyTorch implemented the ndim attribute, making it possible to use matplotlib routines to plot PyTorch tensors directly. I would suspect this works only if tensors are on the CPU, but I haven't checked.

We should figure out which torch version implemented this, bump to that as a minimum, and then simplify our plotting routines. I suspect we could remove many torch2npy calls and just plot tensors directly.