condition no longer needed

BIMSBbioinfo · Mar 4, 2024 · 313ebd2 · 313ebd2
1 parent 64686fd
commit 313ebd2
Showing 1 changed file with 7 additions and 18 deletions.
diff --git a/flexynesis/models/direct_pred_gcnn.py b/flexynesis/models/direct_pred_gcnn.py
@@ -61,28 +61,17 @@ def __init__(
         # NOTE: For now we use matrices, so number of node input features is 1.
         input_dims = [1 for _ in range(len(layers))]
 
-        if self.gnn_conv_type == 'GCNN':
-            self.encoders = nn.ModuleList(
-                [
-                    GCNN(
+        self.encoders = nn.ModuleList(
+            [
+                GNNs(
                         input_dim=input_dims[i],
                         hidden_dim=int(self.config["hidden_dim"]),  # int because of pyg
                         output_dim=self.config["latent_dim"],
+                        act = self.config['activation'],
+                        conv = self.gnn_conv_type
                     )
-                    for i in range(len(layers))
-                ])
-        else: 
-            self.encoders = nn.ModuleList(
-                [
-                    GNNs(
-                            input_dim=input_dims[i],
-                            hidden_dim=int(self.config["hidden_dim"]),  # int because of pyg
-                            output_dim=self.config["latent_dim"],
-                            act = self.config['activation'],
-                            conv = self.gnn_conv_type
-                        )
-                        for i in range(len(layers))       
-                ])
+                    for i in range(len(layers))       
+            ])
 
         # Init output layers
         self.MLPs = nn.ModuleDict()