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.
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
torch2npycalls and just plot tensors directly.