diff --git a/project/01_1_train_VAE.ipynb b/project/01_1_train_VAE.ipynb
index b5d48ba77..a100b8b74 100644
--- a/project/01_1_train_VAE.ipynb
+++ b/project/01_1_train_VAE.ipynb
@@ -445,7 +445,7 @@
     "    val_df=data.val_y,\n",
     "    model=models.vae.VAE,\n",
     "    model_kwargs=dict(n_features=data.train_X.shape[-1],\n",
-    "                      h_layers=args.hidden_layers,\n",
+    "                      n_neurons=args.hidden_layers,\n",
     "                      # last_encoder_activation=None,\n",
     "                      last_decoder_activation=None,\n",
     "                      dim_latent=args.latent_dim),\n",
diff --git a/project/01_1_train_VAE.py b/project/01_1_train_VAE.py
index e5f91ff59..4943c9d0f 100644
--- a/project/01_1_train_VAE.py
+++ b/project/01_1_train_VAE.py
@@ -6,7 +6,7 @@
 #       extension: .py
 #       format_name: percent
 #       format_version: '1.3'
-#       jupytext_version: 1.14.5
+#       jupytext_version: 1.15.0
 #   kernelspec:
 #     display_name: Python 3
 #     language: python
@@ -233,7 +233,7 @@
     val_df=data.val_y,
     model=models.vae.VAE,
     model_kwargs=dict(n_features=data.train_X.shape[-1],
-                      h_layers=args.hidden_layers,
+                      n_neurons=args.hidden_layers,
                       # last_encoder_activation=None,
                       last_decoder_activation=None,
                       dim_latent=args.latent_dim),
diff --git a/vaep/models/ae.py b/vaep/models/ae.py
index c3a8cd8e6..29cc99dff 100644
--- a/vaep/models/ae.py
+++ b/vaep/models/ae.py
@@ -70,7 +70,7 @@ def get_preds_from_df(df: pd.DataFrame,
 
 
 
-
+leaky_relu_default = nn.LeakyReLU(.1)
 
 class Autoencoder(nn.Module):
     """Autoencoder base class.
@@ -79,8 +79,8 @@ class Autoencoder(nn.Module):
 
     def __init__(self,
                  n_features: int,
-                 n_neurons: Union[int, list],
-                 activation=nn.LeakyReLU(.1),
+                 n_neurons: Union[int, List[int]],
+                 activation=leaky_relu_default,
                  last_decoder_activation=None,
                  dim_latent: int = 10):
         """Initialize an Autoencoder
@@ -106,7 +106,7 @@ def __init__(self,
         self.layers = [n_features, *self.n_neurons]
         self.dim_latent = dim_latent
 
-        #define architecture hidden layer
+        # define architecture hidden layer
         def build_layer(in_feat, out_feat):
             return [nn.Linear(in_feat, out_feat),
                     nn.Dropout(0.2),
@@ -192,15 +192,15 @@ def get_missing_values(df_train_wide: pd.DataFrame,
 
 #     def __init__(self,
 #                  n_features: int,
-#                  h_layers: int,
+#                  n_neurons: int,
 #                  activation=nn.LeakyReLU,
 #                  last_encoder_activation=nn.LeakyReLU,
 #                  last_decoder_activation=None,
 #                  dim_latent: int = 10):
 
 #         super().__init__()
-#         self.n_features, self.h_layers = n_features, list(L(h_layers))
-#         self.layers = [n_features, *self.h_layers]
+#         self.n_features, self.n_neurons = n_features, list(L(n_neurons))
+#         self.layers = [n_features, *self.n_neurons]
 #         self.dim_latent = dim_latent
 
 #         # Encoder
diff --git a/vaep/models/vae.py b/vaep/models/vae.py
index 79436acd1..e7e2e8401 100644
--- a/vaep/models/vae.py
+++ b/vaep/models/vae.py
@@ -13,23 +13,23 @@
 from torch import nn
 import torch.nn.functional as F
 
-
+leaky_relu_default = nn.LeakyReLU(.1)
 
 class VAE(nn.Module):
     def __init__(self,
                  n_features: int,
-                 h_layers: List[str],
-                 activation=nn.LeakyReLU(.1),
-                #  last_encoder_activation=nn.LeakyReLU(.1),
+                 n_neurons: List[int],
+                 activation=leaky_relu_default,
+                 #  last_encoder_activation=leaky_relu_default,
                  last_decoder_activation=None,
                  dim_latent: int = 10):
         super().__init__()
-        #set up hyperparameters
-        self.n_features, self.h_layers = n_features, list(h_layers)
-        self.layers = [n_features, *self.h_layers]
+        # set up hyperparameters
+        self.n_features, self.n_neurons = n_features, list(n_neurons)
+        self.layers = [n_features, *self.n_neurons]
         self.dim_latent = dim_latent
 
-        #define architecture hidden layer
+        # define architecture hidden layer
         def build_layer(in_feat, out_feat):
             return [nn.Linear(in_feat, out_feat),
                     nn.Dropout(0.2),