typo

pyg-team · Dec 31, 2021 · 444177a · 444177a
1 parent b5963a5
commit 444177a
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 6 deletions.
diff --git a/graphgym/custom_graphgym/loader/example.py b/graphgym/custom_graphgym/loader/example.py
@@ -1,8 +1,5 @@
 from torch_geometric.datasets import QM7b
 
-from torch_geometric.graphgym.register import register_loader
+from torch_geometric.graphgym.register import register_dataset
 
-
-@register_loader('qm7b')
-def qm7b_dataset(root_dir: str):
-    return QM7b(root_dir)
+register_dataset('qm7b', QM7b)
diff --git a/test/graphgym/test_graphgym.py b/test/graphgym/test_graphgym.py
@@ -31,7 +31,6 @@ def test_run_single_graphgym():
     cfg.dataset.dir = osp.join('/', 'tmp', str(random.randrange(sys.maxsize)))
     dump_cfg(cfg)
     set_printing()
-    return
 
     seed_everything(cfg.seed)
     auto_select_device()