diff --git a/setup.py b/setup.py
index c0a508e8fe..fe7f0c0144 100644
--- a/setup.py
+++ b/setup.py
@@ -30,6 +30,7 @@
     "matplotlib",  # only needed for tensorboard export
     "msgpack",
     "optax",
+    "tensorstore",
     "rich>=11.1",
     "typing_extensions>=4.1.1",
     "PyYAML>=5.4.1",