diff --git a/requirements/tpu.txt b/requirements/tpu.txt
index e8e3b0af5db8c..e071c604b5c0b 100644
--- a/requirements/tpu.txt
+++ b/requirements/tpu.txt
@@ -9,6 +9,7 @@ setuptools-scm>=8
 wheel
 jinja2
 ray[default]
+ray[data]
 
 # Install torch_xla
 --pre