Disable flax migration to orbax for now (#107)

awslabs · Jul 26, 2023 · 1e8b2f8 · 1e8b2f8
1 parent e63482e
commit 1e8b2f8
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 1 deletion.
diff --git 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
@@ -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 <[email protected]>", "Alberto Gasparin <[email protected]>"]
 license = "Apache-2.0"