tensorboard jax[cpu] # pip install --upgrade "jax[cuda]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html # conda install cuda -c nvidia
