diff --git a/poetry.lock b/poetry.lock
index 1d661e1..deaf79a 100644
--- a/poetry.lock
+++ b/poetry.lock
@@ -888,7 +888,6 @@ jax = ">=0.4.2"
 msgpack = "*"
 numpy = ">=1.12"
 optax = "*"
-orbax = "*"
 PyYAML = ">=5.4.1"
 rich = ">=11.1"
 tensorstore = "*"
@@ -2143,33 +2142,6 @@ jaxlib = ">=0.1.37"
 numpy = ">=1.18.0"
 typing-extensions = ">=3.10.0"
 
-[[package]]
-name = "orbax"
-version = "0.1.3"
-description = "Orbax"
-category = "main"
-optional = false
-python-versions = ">=3.8"
-files = [
-    {file = "orbax-0.1.3-py3-none-any.whl", hash = "sha256:a5268b372f28dd085ef1c400f0eb7514fd235c719989b5caa145f1ae5773811f"},
-    {file = "orbax-0.1.3.tar.gz", hash = "sha256:e65171df1734da28478d8a2f2da837416e7c7a026a6280dc93264324ea40df5a"},
-]
-
-[package.dependencies]
-absl-py = "*"
-cached_property = "*"
-etils = "*"
-flax = "*"
-importlib_resources = "*"
-jax = ">=0.4.1"
-jaxlib = "*"
-numpy = "*"
-pyyaml = "*"
-tensorstore = ">=0.1.20"
-
-[package.extras]
-dev = ["pytest", "pytest-xdist"]
-
 [[package]]
 name = "packaging"
 version = "23.0"
diff --git a/pyproject.toml b/pyproject.toml
index fd2f963..1727b31 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -16,7 +16,6 @@ einop = "*"
 numpy = "*" 
 # jax-metrics = {path = "../jax_metrics"}
 jax-metrics = "^0.2.1"
-orbax = "<0.1.4"
 
 [tool.poetry.group.test.dependencies]
 pytest = ">=7.1.3"