turn on mlps for batch variables

BIMSBbioinfo · Nov 23, 2023 · 81ccf31 · 81ccf31
1 parent 4a1a930
commit 81ccf31
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 3 deletions.
diff --git a/flexynesis/models/direct_pred.py b/flexynesis/models/direct_pred.py
@@ -37,7 +37,7 @@ def __init__(self, config, dataset, target_variables, batch_variables = None, va
                 output_dim=self.config['latent_dim']) for i in range(len(layers))])
 
         self.MLPs = nn.ModuleDict()  # using ModuleDict to store multiple MLPs
-        for var in self.target_variables:
+        for var in self.variables:
             if self.dataset.variable_types[var] == 'numerical':
                 num_class = 1
             else:

diff --git a/flexynesis/models/supervised_vae.py b/flexynesis/models/supervised_vae.py
@@ -69,7 +69,7 @@ def __init__(self,  config, dataset, target_variables, batch_variables = None, v
         # define supervisor heads
         # using ModuleDict to store multiple MLPs
         self.MLPs = nn.ModuleDict()         
-        for var in self.target_variables:
+        for var in self.variables:
             if self.dataset.variable_types[var] == 'numerical':
                 num_class = 1
             else:

diff --git a/flexynesis/models/triplet_encoder.py b/flexynesis/models/triplet_encoder.py
@@ -97,7 +97,7 @@ def __init__(self, config, dataset, target_variables, batch_variables = None, va
 
         # define supervisor heads for both target and batch variables 
         self.MLPs = nn.ModuleDict() # using ModuleDict to store multiple MLPs
-        for var in self.target_variables:
+        for var in self.variables:
             if self.variable_types[var] == 'numerical':
                 num_class = 1
             else: