diff --git a/graphgym/loader_pyg.py b/graphgym/loader_pyg.py
index 7b7e6735..01bf7124 100644
--- a/graphgym/loader_pyg.py
+++ b/graphgym/loader_pyg.py
@@ -203,11 +203,11 @@ def set_dataset_info(dataset):
 
     # count number of dataset splits
     cfg.share.num_splits = 1
-    for key in dataset.data.keys:
+    for key in dataset.data.keys():
         if 'val' in key:
             cfg.share.num_splits += 1
             break
-    for key in dataset.data.keys:
+    for key in dataset.data.keys():
         if 'test' in key:
             cfg.share.num_splits += 1
             break