diff --git a/examples/graphproppred/mol/main_pyg.py b/examples/graphproppred/mol/main_pyg.py
index 30372f83..58e1466c 100644
--- a/examples/graphproppred/mol/main_pyg.py
+++ b/examples/graphproppred/mol/main_pyg.py
@@ -11,10 +11,7 @@
 
 ### importing OGB
 
-### for loading dataset
-from ogb.graphproppred.dataset_pyg import PygGraphPropPredDataset
-### for evaluation
-from ogb.graphproppred import Evaluator
+from ogb.graphproppred import PygGraphPropPredDataset, Evaluator
 
 cls_criterion = torch.nn.BCEWithLogitsLoss()
 reg_criterion = torch.nn.MSELoss()
@@ -166,4 +163,4 @@ def main():
 
 
 if __name__ == "__main__":
-    main()
\ No newline at end of file
+    main()