We read every piece of feedback, and take your input very seriously.
To see all available qualifiers, see our documentation.
There was an error while loading. Please reload this page.
1 parent 736266e commit c3348a5Copy full SHA for c3348a5
1 file changed
community/examples/gke-tpu-v6/tpu-multislice.yaml
@@ -41,7 +41,7 @@ spec:
41
cloud.google.com/gke-tpu-topology: 4x4
42
containers:
43
- name: jax-tpu
44
- image: python:3.10
+ image: us-docker.pkg.dev/cloud-tpu-images/jax-ai-image/tpu:latest
45
ports:
46
- containerPort: 8471
47
- containerPort: 8080
@@ -52,7 +52,6 @@ spec:
52
- bash
53
- -c
54
- |
55
- pip install "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
56
python -c 'import jax; print("Global device count:", jax.device_count())'
57
sleep 60
58
resources:
0 commit comments