added more flow models

flatironinstitute · Sep 11, 2024 · 70e5a01 · 70e5a01
1 parent 7ffcba3
commit 70e5a01
Showing 1 changed file with 7 additions and 1 deletion.
diff --git a/src/cryo_sbi/inference/models/build_models.py b/src/cryo_sbi/inference/models/build_models.py
@@ -21,10 +21,16 @@ def build_npe_flow_model(config: dict, **embedding_kwargs) -> nn.Module:
 
     if config["MODEL"] == "MAF":
         model = zuko.flows.MAF
+    elif config["MODEL"] == "GF":
+        model = zuko.flows.GF
+    elif config["MODEL"] == "CNF":
+        model = zuko.flows.CNF
+    elif config["MODEL"] == "UMNN":
+        model = zuko.flows. UMNN
     elif config["MODEL"] == "NSF":
         model = zuko.flows.NSF
     elif config["MODEL"] == "SOSPF":
-        model = partial(zuko.flows.SOSPF, polynomials=16, degree=5)
+        model = partial(zuko.flows.SOSPF, polynomials=8, degree=5)
     else:
         raise NotImplementedError(
             f"Model : {config['MODEL']} has not been implemented yet!"