removed GELU function from current estimator

flatironinstitute · Aug 11, 2023 · 418523c · 418523c
1 parent 625f927
commit 418523c
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/cryo_sbi/inference/models/build_models.py b/src/cryo_sbi/inference/models/build_models.py
@@ -49,7 +49,7 @@ def build_npe_flow_model(config: dict, **embedding_kwargs) -> nn.Module:
         flow=model,
         theta_shift=config["THETA_SHIFT"],
         theta_scale=config["THETA_SCALE"],
-        **{"activation": nn.GELU},#partial(nn.LeakyReLU, 0.1), nn.GELU
+        **{"activation": partial(nn.LeakyReLU, 0.1)}, # TDOD add changeable activation #partial(nn.LeakyReLU, 0.1), nn.GELU
     )
     print("Training with GELU")