From 752521d6bd6270ac1783822fd182bca01478d71a Mon Sep 17 00:00:00 2001
From: kaminow <benjamin.kaminow@choderalab.org>
Date: Wed, 1 Nov 2023 14:56:32 -0400
Subject: [PATCH] Update LigandOnlyModel return to match other Model classes.

---
 mtenn/model.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/mtenn/model.py b/mtenn/model.py
index 4817091..f784097 100644
--- 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]