ae elbow work

Deep-SAD-PyTorch/src/networks/main.py

@@ -9,7 +9,7 @@ from .subter_LeNet_Split import SubTer_LeNet_Split, SubTer_LeNet_Split_Autoencod
 from .vae import VariationalAutoencoder
 
 
-def build_network(net_name, ae_net=None):
+def build_network(net_name, ae_net=None, rep_dim=1024):
     """Builds the neural network."""
 
     implemented_networks = (
@@ -46,7 +46,7 @@ def build_network(net_name, ae_net=None):
         net = MNIST_LeNet()
 
     if net_name == "subter_LeNet":
-        net = SubTer_LeNet()
+        net = SubTer_LeNet(rep_dim=rep_dim)
 
     if net_name == "subter_LeNet_Split":
         net = SubTer_LeNet_Split()