diff --git a/autogl/datasets/_ogb.py b/autogl/datasets/_ogb.py
index 7192a0c..0679b30 100644
--- a/autogl/datasets/_ogb.py
+++ b/autogl/datasets/_ogb.py
@@ -39,7 +39,7 @@ class _OGBNDatasetUtil(_OGBDatasetUtil):
         edge_index = SparseTensor(row=torch.tensor(edge_index[0]), col=torch.tensor(edge_index[1]), value=edge_feat, sparse_sizes=(num_nodes, num_nodes))
         edge_index = edge_index.to_symmetric()
         row, col, _ = edge_index.coo()
-        edge_index = [row.cpu().detach().numpy(), col.cpu().detach().numpy()]
+        edge_index = np.array([row.cpu().detach().numpy(), col.cpu().detach().numpy()])
         homogeneous_static_graph: GeneralStaticGraph = (
             GeneralStaticGraphGenerator.create_homogeneous_static_graph(
                 dict([
diff --git a/examples/nodeclf_ogb.py b/examples/nodeclf_ogb.py
index 3fbac6a..bb8a34e 100644
--- a/examples/nodeclf_ogb.py
+++ b/examples/nodeclf_ogb.py
@@ -174,8 +174,8 @@ def main():
     num_classes = len(np.unique(labels.numpy()))
 
     if args.use_sage:
-        model = SAGE(data.num_features, args.hidden_channels,
-                     dataset.num_classes, args.num_layers,
+        model = SAGE(dataset[0].nodes.data[feat].size(1), args.hidden_channels,
+                     num_classes, args.num_layers,
                      args.dropout).to(device)
     else:
         model = GCN(dataset[0].nodes.data[feat].size(1), args.hidden_channels,