Skip to content

Commit

Permalink
[Chem Model Zoo] Hot fix for README (dmlc#1071)
Browse files Browse the repository at this point in the history
* Fix outdated example

* Hot fix
  • Loading branch information
mufeili authored Dec 5, 2019
1 parent 7784f06 commit bc4f435
Showing 1 changed file with 3 additions and 3 deletions.
6 changes: 3 additions & 3 deletions examples/pytorch/model_zoo/chem/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -63,16 +63,16 @@ as front end and Set2Set for output prediction.
### Example Usage of Pre-trained Models

```python
from dgl.data.chem import Tox21
from dgl.data.chem import Tox21, smiles_to_bigraph, CanonicalAtomFeaturizer
from dgl import model_zoo

dataset = Tox21()
dataset = Tox21(smiles_to_bigraph, CanonicalAtomFeaturizer())
model = model_zoo.chem.load_pretrained('GCN_Tox21') # Pretrained model loaded
model.eval()

smiles, g, label, mask = dataset[0]
feats = g.ndata.pop('h')
label_pred = model(feats, g)
label_pred = model(g, feats)
print(smiles) # CCOc1ccc2nc(S(N)(=O)=O)sc2c1
print(label_pred[:, mask != 0]) # Mask non-existing labels
# tensor([[-0.7956, 0.4054, 0.4288, -0.5565, -0.0911,
Expand Down

0 comments on commit bc4f435

Please sign in to comment.