fix for transformer conv inference and doc string

src/anemoi/models/layers/block.py

@@ -600,28 +600,7 @@ def forward(
             ), "Only batch size of 1 is supported when model is sharded across GPUs"
 
         query, key, value, edges = self.shard_qkve_heads(query, key, value, edges, shapes, batch_size, model_comm_group)
-
-        # TODO: Is this alright?
-        num_chunks = self.num_chunks if self.training else 4  # reduce memory for inference
-
-        if num_chunks > 1:
-            edge_index_list = torch.tensor_split(edge_index, num_chunks, dim=1)
-            edge_attr_list = torch.tensor_split(edges, num_chunks, dim=0)
-            for i in range(num_chunks):
-                out1 = self.conv(
-                    query=query,
-                    key=key,
-                    value=value,
-                    edge_attr=edge_attr_list[i],
-                    edge_index=edge_index_list[i],
-                    size=size,
-                )
-                if i == 0:
-                    out = torch.zeros_like(out1)
-                out = out + out1
-        else:
-            out = self.conv(query=query, key=key, value=value, edge_attr=edges, edge_index=edge_index, size=size)
-
+        out = self.conv(query=query, key=key, value=value, edge_attr=edges, edge_index=edge_index, size=size)
         out = self.shard_output_seq(out, shapes, batch_size, model_comm_group)
         out = self.projection(out + x_r)
 

src/anemoi/models/layers/conv.py

@@ -7,7 +7,6 @@
 # nor does it submit to any jurisdiction.
 #
 
-import math
 from typing import Optional
 
 import torch
@@ -77,7 +76,11 @@ def aggregate(self, edges_new: Tensor, edge_index: Adj, dim_size: Optional[int]
 
 
 class GraphTransformerConv(MessagePassing):
-    """Message passing part of graph transformer operator."""
+    """Message passing part of graph transformer operator. 
+
+    Adapted from 'Masked Label Prediction: Unified Message Passing Model for Semi-Supervised Classification'
+    (https://arxiv.org/abs/2009.03509)
+    """
 
     def __init__(
         self,
@@ -130,7 +133,7 @@ def message(
         if edge_attr is not None:
             key_j = key_j + edge_attr
 
-        alpha = (query_i * key_j).sum(dim=-1) / math.sqrt(self.out_channels)
+        alpha = (query_i * key_j).sum(dim=-1) / self.out_channels ** 0.5
 
         alpha = softmax(alpha, index, ptr, size_i)
         alpha = dropout(alpha, p=self.dropout, training=self.training)