diff --git a/pyproject.toml b/pyproject.toml
index 49c4d06..4590269 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -18,6 +18,7 @@ dependencies = [
     "orbax-checkpoint",
     "uvicorn",
     "requests",
+    "packaging",
 ]
 
 # `version` is automatically set by flit to use `my_project.__version__`
diff --git a/requirements.txt b/requirements.txt
index 1d0cd4a..0320364 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,6 +1,7 @@
-jax[cpu]>=0.5.1
 absl-py
+fastapi
+jax>=0.7.2
 orbax-checkpoint
 uvicorn
-fastapi
-google-cloud-logging
+requests
+packaging