Update tpu-test.py

GoogleCloudPlatform · Mar 6, 2024 · f8b6687 · f8b6687
1 parent b146cb7
commit f8b6687
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/applications/ray/kuberay-tpu-webhook/tests/tpu-test.py b/applications/ray/kuberay-tpu-webhook/tests/tpu-test.py
@@ -1,6 +1,6 @@
 import ray
 
-@ray.remote(resources={"google.com/tpu": 4})
+@ray.remote(resources={"TPU": 4})
 def tpu_cores():
     import jax
 
@@ -23,4 +23,4 @@ def tpu_cores():
 
 num_workers = 4
 result = [tpu_cores.remote() for _ in range(num_workers)]
-print(ray.get(result))
+print(ray.get(result))