diff --git a/models/UMNN/UMNNMAF.py b/models/UMNN/UMNNMAF.py
index 336c9c2..4efbd42 100644
--- a/models/UMNN/UMNNMAF.py
+++ b/models/UMNN/UMNNMAF.py
@@ -6,7 +6,7 @@
 import math
 from .made import MADE, ConditionnalMADE
 
-dict_act_func = {"Sigmoid": nn.Sigmoid(), "ELU": nn.ELU()}
+dict_act_func = {"Sigmoid": nn.Sigmoid(), "ELU": lambda x: nn.ELU(x) + 1.}
 
 def _flatten(sequence):
     flat = [p.contiguous().view(-1) for p in sequence]
@@ -219,7 +219,7 @@ def force_lipschitz(self, L=1.5):
 
 class EmbeddingNetwork(nn.Module):
     def __init__(self, in_d, hiddens_embedding=[50, 50, 50, 50], hiddens_integrand=[50, 50, 50, 50], out_made=1,
-                 cond_in=0, device="cpu"):
+                 cond_in=0, act_func='ELU', device="cpu"):
         super().__init__()
         self.m_embeding = None
         self.device = device
@@ -229,7 +229,7 @@ def __init__(self, in_d, hiddens_embedding=[50, 50, 50, 50], hiddens_integrand=[
                                          natural_ordering=True).to(device)
         else:
             self.made = MADE(in_d, hiddens_embedding, in_d * (out_made), num_masks=1, natural_ordering=True).to(device)
-        self.parallel_nets = IntegrandNetwork(in_d, 1 + out_made, hiddens_integrand, 1, device=device)
+        self.parallel_nets = IntegrandNetwork(in_d, 1 + out_made, hiddens_integrand, 1, act_func=act_func, device=device)
 
     def make_embeding(self, x_made, context=None):
         self.m_embeding = self.made.forward(x_made, context)
diff --git a/models/UMNN/UMNNMAFFlow.py b/models/UMNN/UMNNMAFFlow.py
index a9d2976..2eae708 100644
--- a/models/UMNN/UMNNMAFFlow.py
+++ b/models/UMNN/UMNNMAFFlow.py
@@ -38,7 +38,7 @@ def __getitem__(self, i):
 class UMNNMAFFlow(nn.Module):
 
     def __init__(self, nb_flow=1, nb_in=1, hidden_derivative=[50, 50, 50, 50], hidden_embedding=[50, 50, 50, 50],
-                 embedding_s=20, nb_steps=50, solver="CC", cond_in=0, device="cpu"):
+                 embedding_s=20, nb_steps=50, act_func='ELU', solver="CC", cond_in=0, device="cpu"):
         """
         UMNNMAFFlow class is a normalizing flow made of UMNNMAF blocks.
         :int nb_flow: The number of components in the flow
@@ -57,7 +57,7 @@ def __init__(self, nb_flow=1, nb_in=1, hidden_derivative=[50, 50, 50, 50], hidde
 
         self.nets = ListModule(self, "Flow")
         for i in range(nb_flow):
-            auto_net = EmbeddingNetwork(nb_in, hidden_embedding, hidden_derivative, embedding_s,
+            auto_net = EmbeddingNetwork(nb_in, hidden_embedding, hidden_derivative, embedding_s, act_func='ELU',
                                                      device=device, cond_in=cond_in).to(device)
 
             model = UMNNMAF(auto_net, nb_in, nb_steps, device, solver=solver).to(device)
diff --git a/models/UMNN/__init__.py b/models/UMNN/__init__.py
index eec7f44..9d7224c 100644
--- a/models/UMNN/__init__.py
+++ b/models/UMNN/__init__.py
@@ -1,2 +1,3 @@
 from .UMNNMAFFlow import UMNNMAFFlow
-from .MonotonicNN import MonotonicNN, IntegrandNN
\ No newline at end of file
+from .MonotonicNN import MonotonicNN, IntegrandNN
+from .made import MADE
\ No newline at end of file
diff --git a/models/__init__.py b/models/__init__.py
index c7722fe..227ee63 100644
--- a/models/__init__.py
+++ b/models/__init__.py
@@ -1 +1 @@
-from models.UMNN import UMNNMAFFlow
+from models.UMNN import UMNNMAFFlow, MADE