add flax to setup

e3nn · Dec 14, 2022 · 1a70584 · 1a70584
1 parent 24bff4c
commit 1a70584
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/setup.py b/setup.py
@@ -35,6 +35,7 @@ def find_version(*file_paths):
     install_requires=[
         "jax",
         "dm-haiku",
+        "flax",
         "optax",
         "sympy",
         "numpy",