diff --git a/constellation/ConstellationNet.py b/constellation/ConstellationNet.py index 39fe417..a2b10c5 100644 --- a/constellation/ConstellationNet.py +++ b/constellation/ConstellationNet.py @@ -42,6 +42,8 @@ class ConstellationNet(nn.Module): for layer_size in encoder_layers_sizes: encoder_layers.append(nn.Linear(prev_layer_size, layer_size)) + encoder_layers.append(nn.ReLU()) + encoder_layers.append(nn.BatchNorm1d(layer_size)) prev_layer_size = layer_size encoder_layers += [ @@ -61,6 +63,8 @@ class ConstellationNet(nn.Module): for layer_size in decoder_layers_sizes: decoder_layers.append(nn.Linear(prev_layer_size, layer_size)) + encoder_layers.append(nn.ReLU()) + decoder_layers.append(nn.BatchNorm1d(layer_size)) prev_layer_size = layer_size # Softmax is not used at the end of the network because the