From 3597a20feaf744b4c5c381233065e923a3f8bf79 Mon Sep 17 00:00:00 2001
From: Alex Cabrera <alex.cabrera@gmail.com>
Date: Mon, 19 Aug 2024 17:47:40 -0700
Subject: [PATCH] correctly set base model to eval mode

---
 molfeat/trans/pretrained/dgl_pretrained.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/molfeat/trans/pretrained/dgl_pretrained.py b/molfeat/trans/pretrained/dgl_pretrained.py
index 41eb630..8c6a7bc 100644
--- a/molfeat/trans/pretrained/dgl_pretrained.py
+++ b/molfeat/trans/pretrained/dgl_pretrained.py
@@ -74,8 +74,8 @@ def from_pretrained(cls, model_name: str):
         import dgllife
 
         base_model = dgllife.model.load_pretrained(model_name)
+        base_model.eval()
         model = DGLModel(name=model_name)
-        model.eval()
         model._model = base_model
         return model