Skip to content

Commit

Permalink
correctly set base model to eval mode
Browse files Browse the repository at this point in the history
  • Loading branch information
cabreraalex authored Aug 20, 2024
1 parent 40414ae commit 3597a20
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 3597a20

Please sign in to comment.