Skip to content

Commit 7b74d97

Browse files
committed
update supported datasets
1 parent bb4ee4d commit 7b74d97

File tree

1 file changed

+7
-1
lines changed

1 file changed

+7
-1
lines changed

graphgallery/data/planetoid.py

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,10 +21,16 @@ class Planetoid(Dataset):
2121

2222
github_url = "https://raw.githubusercontent.com/EdisonLeeeee/GraphData/master/datasets/planetoid"
2323
obj_names = ["adj.pkl", "feature.npy", "label.npy", "idx_train.npy", "idx_val.npy", "idx_test.npy"]
24+
supported_datasets = ('citeseer', 'cora', 'pubmed')
2425

2526
def __init__(self, name, root=None):
27+
name = name.lower()
28+
29+
if not name in self.supported_datasets:
30+
raise ValueError(f"Currently only support for these datasets {self.supported_datasets}.")
31+
2632
super().__init__(name, root)
27-
33+
2834
self.download_dir = osp.join(self.root, 'planetoid', name, 'raw')
2935
self.processed_dir = osp.join(self.root, 'planetoid', name, 'processed')
3036

0 commit comments

Comments
 (0)