RL-2070: remove activation after last layer

(Issue RL-2070 - Implement PyG gnn block)

Author: EnliteAI Bot

maze/perception/blocks/feed_forward/gnn_pyg.py

@@ -63,7 +63,7 @@ def forward(self, x: torch.Tensor, edge_index: torch.Tensor, edge_attr: torch.Te
 
         :param x: Note feature matrix, [n_nodes, in_features] or [B, n_nodes, in_features].
         :param edge_index: The graph connectivity in COO format, [2, E] or [B, 2, E].
-        :param edge_attr: The edge attributes, [E, D] or [B, E, D]. D is the dimension of edge attribute.
+        :param edge_attr: The edge attributes, [E, D] or [B, E, D]. D is the edge attribute dimension.
         :return: Output tensor.
         """
         if x.dim() == 2:  # Single graph (no batch)
@@ -192,7 +192,9 @@ def build_layer_dict(self) -> OrderedDict:
                 bias=self.bias,
                 gnn_type=self.gnn_type,
             )
-            layer_dict[f'activation_{layer_idx}_{self.non_lin.__name__}'] = self.non_lin()
+            # Add activation function only for intermediate layers
+            if layer_idx < len(self.hidden_features) - 1:
+                layer_dict[f'activation_{layer_idx}_{self.non_lin.__name__}'] = self.non_lin()
             in_feats = out_feats
 
         return layer_dict

maze/test/perception/blocks/test_gnn_pyg.py

@@ -45,7 +45,7 @@ def test_gnn_layer_forward(gnn_type, batch_size, bias):
     if batch_size is None:
         x = torch.randn(n_nodes, in_features)
         edge_index = torch.tensor([[0, 1, 2], [1, 2, 0]], dtype=torch.long)
-        edge_attr = torch.randn(edge_index.shape[1])  # (E,) for GCN or even random
+        edge_attr = torch.randn(edge_index.shape[1])  # (E,) for GCN
     else:
         x = torch.randn(batch_size, n_nodes, in_features)
         edge_index = []
@@ -70,7 +70,7 @@ def test_gnn_layer_forward(gnn_type, batch_size, bias):
 @pytest.mark.parametrize("gnn_type", SUPPORTED_GNNS)
 def test_gnn_block_forward(gnn_type: str):
     """
-    Test the forward pass of GNNBlockPyG with dummy data.
+    Test the forward pass of GNNBlockPyG.
     """
 
     batch_size = 2
@@ -91,7 +91,6 @@ def test_gnn_block_forward(gnn_type: str):
 
     node_feats = torch.randn(batch_size, n_nodes, in_features)
 
-    # Edge index shape expected: (batch_size, 2, E)  =>  E = 10 here
     edge_index = []
     edge_attr = []
     for _ in range(batch_size):