Convert to Pytorch

[Json-To-String]

I'd love to ditch Tensorflow and just use PyTorch + Tensorboard. Could always support both as backends.

[nickvazz]

https://github.com/jax-ml/jax is also another good option

[Json-To-String]

Hmm, this could be a great excuse to learn Jax. Between it and pytorch are there great arguments for one over the other? My current understanding is that its an ease of use difference as well as Jax being a bit more drop-in for numpy?