From 4e6b640de83b803cd1ef526029f78a239a740114 Mon Sep 17 00:00:00 2001 From: DEGIACOMI Date: Fri, 23 Feb 2024 21:08:34 +0000 Subject: [PATCH] bugfix for foldingnet foldingnet can now be correctly used with more than 2D latent space --- src/molearn/models/foldingnet.py | 2 +- src/molearn/models/small_foldingnet.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/src/molearn/models/foldingnet.py b/src/molearn/models/foldingnet.py index a25e613..d60c03e 100644 --- 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 index c31e837..4684a39 100644 --- 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)