From 38e6a174cdeee9f67aef3dc246ac10896627005f Mon Sep 17 00:00:00 2001 From: Jesper Dramsch Date: Mon, 10 Jun 2024 15:19:39 +0200 Subject: [PATCH] Simplify exponent --- src/anemoi/models/layers/conv.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/anemoi/models/layers/conv.py b/src/anemoi/models/layers/conv.py index 681018f..f71683e 100644 --- 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)