diff --git a/fortuna/__init__.py b/fortuna/__init__.py
index e69de29b..e28a435b 100755
--- a/fortuna/__init__.py
+++ b/fortuna/__init__.py
@@ -0,0 +1,3 @@
+import flax
+
+flax.config.update("flax_use_orbax_checkpointing", False)
diff --git a/pyproject.toml b/pyproject.toml
index 1720bfcb..1df24d24 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "aws-fortuna"
-version = "0.1.19"
+version = "0.1.20"
 description = "A Library for Uncertainty Quantification."
 authors = ["Gianluca Detommaso <gianluca.detommaso@gmail.com>", "Alberto Gasparin <albgas@amazon.de>"]
 license = "Apache-2.0"