why in_length is needed ? if n_atom_properties is available

Author: aditya0by0

chebai_graph/models/_gat.py

@@ -13,7 +13,6 @@ class GATModelWrapper(GraphBaseNet):
     def __init__(self, config: dict, **kwargs):
         super().__init__(**kwargs)
 
-        self._in_length = int(config.pop("in_length"))
         self._hidden_length = int(config.pop("hidden_length"))
         self._dropout_rate = float(config.pop("dropout_rate", 0.1))
         self._n_conv_layers = int(config.pop("n_conv_layers", 3))
@@ -22,10 +21,11 @@ def __init__(self, config: dict, **kwargs):
         self._n_bond_properties = int(config.pop("n_bond_properties", 0))
         self._n_molecule_properties = int(config.pop("n_molecule_properties", 0))
         self._gat = GAT(
-            in_channels=self._in_length,
+            in_channels=self._n_atom_properties,
             hidden_channels=self._hidden_length,
             num_layers=self._n_conv_layers,
             dropout=self._dropout_rate,
+            edge_dim=self._n_bond_properties,
             **config,
         )
 

configs/model/gat.yml

@@ -3,11 +3,10 @@ init_args:
   optimizer_kwargs:
     lr: 1e-3
   config:
-    in_length: 256
     hidden_length: 512
     dropout_rate: 0.1
     n_conv_layers: 3
-    heads: 4  # the number of heads should be divisible by output channels (hidden channels if output channel not given)
+    heads: 8  # the number of heads should be divisible by output channels (hidden channels if output channel not given)
     # v2: True  # -- to use `torch_geometric.nn.conv.GATv2Conv` convolution layers, default is GATConv
     n_linear_layers: 3
     n_atom_properties: 158