bugfix for foldingnet

foldingnet can now be correctly used with more than 2D latent space
Degiacomi-Lab · Feb 23, 2024 · 4e6b640 · 4e6b640
1 parent 96bf989
commit 4e6b640
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/src/molearn/models/foldingnet.py b/src/molearn/models/foldingnet.py
@@ -223,7 +223,7 @@ def forward(self, x):
         '''
         x: (B, C)
         '''
-        x = x.view(-1, 2, 1)
+        x = x.view(-1, self.latent_dimension, 1)
         x = self.layer1(x)
         x = self.layer2(x)
         x = self.layer3(x)

diff --git a/src/molearn/models/small_foldingnet.py b/src/molearn/models/small_foldingnet.py
@@ -68,7 +68,7 @@ def forward(self, x):
         '''
         x: (B, C)
         '''
-        x = x.view(-1, 2, 1)
+        x = x.view(-1, self.latent_dimension, 1)
         x = self.layer1(x)
         x = self.layer2(x)
         x = self.layer3(x)