Merge pull request #316 from michellemli/main

mims-harvard · Sep 24, 2024 · 45cc094 · 45cc094
2 parents 80acefc + bf72829
commit 45cc094
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/tdc/resource/primekg.py b/tdc/resource/primekg.py
@@ -33,9 +33,9 @@ def to_nx(self):
 
         G = nx.Graph()
         for i in self.df.relation.unique():
-            G.add_edges_from(self.df[self.df.relation == i][["x_id",
-                                                             "y_id"]].values,
-                             relation=i)
+            G.add_edges_from(
+                self.df[self.df.relation == i][["x_name", "y_name"]].values,
+                relation=i)
         return G
 
     def get_features(self, feature_type):