Correctly instantiate EgoData in SnapDatasets

pyg-team · Oct 5, 2022 · 6ca2332 · 6ca2332
1 parent f709376
commit 6ca2332
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/torch_geometric/datasets/snap_dataset.py b/torch_geometric/datasets/snap_dataset.py
@@ -110,8 +110,8 @@ def read_ego(files, name):
         edge_index = torch.stack([row, col], dim=0)
 
         edge_index, _ = coalesce(edge_index, None, N, N)
-        data = Data(x=x, edge_index=edge_index, circle=circle,
-                    circle_batch=circle_batch)
+        data = EgoData(x=x, edge_index=edge_index, circle=circle,
+                       circle_batch=circle_batch)
 
         data_list.append(data)