From 418523ca9abc385d4300e99d4e00867beaeedf14 Mon Sep 17 00:00:00 2001
From: Dingel321 <Dingel321@gmail.com>
Date: Fri, 11 Aug 2023 17:22:24 -0400
Subject: [PATCH] removed GELU function from current estimator

---
 src/cryo_sbi/inference/models/build_models.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/cryo_sbi/inference/models/build_models.py b/src/cryo_sbi/inference/models/build_models.py
index 1b0ccc3..9a392fd 100644
--- 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")