diff --git a/src/deep_neurographs/machine_learning/heterograph_models.py b/src/deep_neurographs/machine_learning/heterograph_models.py
index bbd52cf..73af1d7 100644
--- a/src/deep_neurographs/machine_learning/heterograph_models.py
+++ b/src/deep_neurographs/machine_learning/heterograph_models.py
@@ -255,7 +255,7 @@ def forward(self, x_dict, edge_index_dict, edge_attr_dict):
 
 # --- Utils ---
 def reformat_edge_key(key):
-    if type(key) is str
+    if type(key) is str:
         return tuple([rm_non_alphanumeric(s) for s in key.split(",")])
     else:
         return key