From 1e8b2f841815bb28b65edc0d0ef0ff135da308d6 Mon Sep 17 00:00:00 2001
From: Gianluca Detommaso <32386694+gianlucadetommaso@users.noreply.github.com>
Date: Wed, 26 Jul 2023 15:25:38 +0200
Subject: [PATCH] Disable flax migration to orbax for now (#107)

---
 fortuna/__init__.py | 3 +++
 pyproject.toml      | 2 +-
 2 files changed, 4 insertions(+), 1 deletion(-)

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"