fix gat v2 share weights issue

Author: aditya0by0

chebai_graph/models/gat.py

@@ -28,7 +28,9 @@ def __init__(self, config: dict, **kwargs):
         super().__init__(config=config, **kwargs)
         self.heads = int(config["heads"])
         self.v2 = bool(config["v2"])
-        self.share_weights = bool(config.get("share_weights", False))
+        local_kwargs = {}
+        if self.v2:
+            local_kwargs["share_weights"] = bool(config.get("share_weights", False))
         self.activation = ELU()  # Instantiate ELU once for reuse.
         self.gat = GAT(
             in_channels=self.in_channels,
@@ -40,7 +42,7 @@ def __init__(self, config: dict, **kwargs):
             heads=self.heads,
             v2=self.v2,
             act=self.activation,
-            share_weights=self.share_weights,
+            **local_kwargs,
         )
 
     def forward(self, batch: dict) -> torch.Tensor: