bring in prior_model from PyG visnet

Co-authored-by: kaminow <51923685+kaminow@users.noreply.github.com>
choderalab · Jan 31, 2024 · ffa571c · ffa571c
1 parent 48721fb
commit ffa571c
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/mtenn/conversion_utils/visnet.py b/mtenn/conversion_utils/visnet.py
@@ -90,6 +90,9 @@ def forward(self, data):
             x = self.visnet.output_model.pre_reduce(x, v)
             x = x * self.visnet.std
 
+        if self.visnet.prior_model is not None:
+            x = self.visnet.prior_model(x, z)
+
             return x
 
         def _get_representation(self):