adjust activation function config

ecmwf · Feb 28, 2025 · 4f6003d · 4f6003d
1 parent e1ceca2
commit 4f6003d
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 3 deletions.
diff --git a/training/src/anemoi/training/config/model/gnn.yaml b/training/src/anemoi/training/config/model/gnn.yaml
@@ -1,4 +1,8 @@
-activation: GELU
+activation:
+  _target_: nn.torch.GELU
+  _partial_: True
+  #dim: ${model.num_channels} # for GLU functions
+
 num_channels: 512
 cpu_offload: False
 output_mask: null

diff --git a/training/src/anemoi/training/config/model/graphtransformer.yaml b/training/src/anemoi/training/config/model/graphtransformer.yaml
@@ -1,4 +1,7 @@
-activation: GELU
+activation:
+  _target_: nn.torch.GELU
+  _partial_: True
+
 num_channels: 1024
 cpu_offload: False
 output_mask: null

diff --git a/training/src/anemoi/training/config/model/transformer.yaml b/training/src/anemoi/training/config/model/transformer.yaml
@@ -1,4 +1,7 @@
-activation: GELU
+activation:
+  _target_: nn.torch.GELU
+  _partial_: True
+
 num_channels: 1024
 cpu_offload: False
 output_mask: null