支持NVIDIA GPU —— 如何运行docker环境下的Jax环境

项目地址：
https://github.com/NVIDIA/JAX-Toolbox

---

具体的安装命令：

地址：
https://github.com/NVIDIA/JAX-Toolbox/pkgs/container/jax

linux/amd64

docker pull ghcr.io/nvidia/jax:mealkit-2024-01-08@sha256:8d2dd85097b5929dcc08a2ad229724d489865f2e6f5a401cf7c0e532b5e8c47b

linux/arm64

docker pull ghcr.io/nvidia/jax:mealkit-2024-01-08@sha256:c83bdd5b432489772efeaa634717cc5fc9cb25984a8ef3adc03a8a70e60b0954