Merge pull request #31 from choderalab/update-lig-model-return

Update LigandOnlyModel return to match other Model classes.
choderalab · Nov 1, 2023 · 011c066 · 011c066
2 parents 6f6d8e8 + 752521d
commit 011c066
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/mtenn/model.py b/mtenn/model.py
@@ -247,6 +247,6 @@ def forward(self, rep):
         pred = self.get_representation(tmp_rep)
 
         if self.readout:
-            return self.readout(pred)
+            return self.readout(pred), [pred]
         else:
-            return pred
+            return pred, [pred]