Skip to content

JAX 'Hello World' on GKE + A100-80GB

Open in Cloud Shell

This tutorial shows how to run a simple JAX program using NVIDIA GPUs A100-80GB on a GKE cluster

Visit https://cloud.google.com/blog/products/containers-kubernetes/machine-learning-with-jax-on-kubernetes-with-nvidia-gpus to follow the tutorial