Skip to content

Commit

Permalink
Merge pull request #104 from cabreraalex/main
Browse files Browse the repository at this point in the history
correctly set base model to eval mode
  • Loading branch information
maclandrol authored Sep 11, 2024
2 parents ff92ee6 + 6d9eb4d commit d787602
Showing 1 changed file with 1 addition and 1 deletion.
2 changes: 1 addition & 1 deletion molfeat/trans/pretrained/dgl_pretrained.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down

0 comments on commit d787602

Please sign in to comment.