diff --git a/src/models.py b/src/models.py
index 015d29e..000d01b 100644
--- a/src/models.py
+++ b/src/models.py
@@ -98,6 +98,7 @@ def init_conv_layers(self):
         for i in range(self.n_partitions):
             part_conv_layers = []
             for j in range(self.n_conv[i]):
+                # TODO Add possibility of changing convolutional layers
                 conv_layer = [
                     gnn.CGConv(
                         channels=self.conv_size[i],