disable flax migration to orbax for now

awslabs · Jul 26, 2023 · 9f1a738 · 9f1a738
1 parent c5bc94f
commit 9f1a738
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/fortuna/__init__.py b/fortuna/__init__.py
@@ -0,0 +1,2 @@
+import flax
+flax.config.update('flax_use_orbax_checkpointing', False)