From f8b6687c2ffdddda8138ee43de9ee712968099c4 Mon Sep 17 00:00:00 2001
From: Richard Liu <39319471+richardsliu@users.noreply.github.com>
Date: Wed, 6 Mar 2024 13:50:04 -0800
Subject: [PATCH] Update tpu-test.py

---
 applications/ray/kuberay-tpu-webhook/tests/tpu-test.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/applications/ray/kuberay-tpu-webhook/tests/tpu-test.py b/applications/ray/kuberay-tpu-webhook/tests/tpu-test.py
index 3c98d88fd..411413638 100644
--- 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))
\ No newline at end of file
+print(ray.get(result))