Simplify exponent

ecmwf · Jun 10, 2024 · 38e6a17 · 38e6a17
1 parent 5faeafa
commit 38e6a17
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/anemoi/models/layers/conv.py b/src/anemoi/models/layers/conv.py
@@ -131,7 +131,7 @@ def message(
         if edge_attr is not None:
             key_j = key_j + edge_attr
 
-        alpha = (query_i * key_j).sum(dim=-1) / self.out_channels ** (1.0 / 2.0)
+        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)