diff --git a/pyproject.toml b/pyproject.toml
index d887255..2dd080b 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -7,7 +7,7 @@ requires-python = ">=3.8"
 
 readme = "README.md"
 
-version = "0.0.54"
+version = "0.0.55"
 
 dependencies = [
     "jax>=0.4.20",
diff --git a/src/fjformer/__init__.py b/src/fjformer/__init__.py
index 17c28b9..41be7cc 100644
--- a/src/fjformer/__init__.py
+++ b/src/fjformer/__init__.py
@@ -68,7 +68,7 @@
 from . import optimizers as optimizers
 from . import linen as linen
 
-__version__ = "0.0.54"
+__version__ = "0.0.55"
 
 __all__ = (
     # Loss and extra function