diff --git a/heterograph_construction.py b/heterograph_construction.py
index 1a8f114c6e61905286aa3020d3c72e06eb79cb20..b12dbdf297b8f64eb8c9e350cc1bbcaf40e554cf 100644
--- a/heterograph_construction.py
+++ b/heterograph_construction.py
@@ -321,8 +321,7 @@ class DISNETConstructor:
 
         # Add the edges to the graph.
         for edge_t in edges_dict.keys():
-            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.
+            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:
@@ -404,7 +403,7 @@ class DISNETConstructor:
     """
     def decodePredictions(self, list, type, n=50, prepared=False, name=''):
         dir = 'results/' + type + '_' + str(n) + name + '_' + 'table.csv'
-        if (type == 'dis_dru_the'):
+        if type == 'dis_dru_the':
             # Read the data
             h, t = self.getNodeInfo(full=False)
             # Give each node its ID