diff --git a/heterograph_construction.py b/heterograph_construction.py
index 1cbc5c8af454ce3afe77709f4f966c9129479482..a8ee51063f1645ab352905cb1b9df9ff86033e98 100644
--- a/heterograph_construction.py
+++ b/heterograph_construction.py
@@ -213,9 +213,9 @@ class DISNETConstructor:
 
             # Delete from the graph the edges of repo_db
             dis_dru_the = pd.concat([dis_dru_the, dis_dru_the_repoDBAll])
-            dis_dru_the.drop_duplicates(keep=False, inplace=True)
+            dis_dru_the.drop_duplicates(keep=False, inplace=True, ignore_index=True)
 
-            dis_dru_the = dis_dru_the[:50355]
+            dis_dru_the = dis_dru_the[:50355] # Those of RepoDB added for the concatenation are deleted.
 
             dis_dru_the_repoDBAll = (
                 torch.tensor(dis_dru_the_repoDBAll['disNID'].astype(np.int32).to_numpy(), dtype=torch.int32,
@@ -321,7 +321,7 @@ class DISNETConstructor:
 
         # Add the edges to the graph.
         for edge_t in edges_dict.keys():
-            for edge in edges[edge_t]:
+            for edge in edges[edge_t]: # If a relation between a phenotype and a drug is in more than one type. The final type will be the last one.
                 try:
                     G.add_edge(int(edge[0]), int(edge[1]), edge_feature=edge[2], edge_type=edge_t)
                 except IndexError:
diff --git a/topN.py b/topN.py
index 83a9d626ccd20a7d0f7e58f789ad14c979d4977b..fd92fb5fe9256fa5ab32b2a72fbd9c1c3f9810a8 100644
--- a/topN.py
+++ b/topN.py
@@ -26,7 +26,7 @@ def filterPreds(original, pred):
 
     new = []
     for i, elem in enumerate(pred):
-        pred_labels = pred[edge, i]
+        pred_labels = pred[elem]
         head = i
         tail = torch.arange(0, len(pred_labels))  # All tails.
         indexH = ((headsO == head).nonzero(as_tuple=True)[0])  # Index of those heads originally present in the graph.