From 4f6003d472c33a77a87d8345b43d024ae3c94a4c Mon Sep 17 00:00:00 2001 From: chebertpinard Date: Fri, 28 Feb 2025 16:30:22 +0000 Subject: [PATCH] adjust activation function config --- training/src/anemoi/training/config/model/gnn.yaml | 6 +++++- .../src/anemoi/training/config/model/graphtransformer.yaml | 5 ++++- training/src/anemoi/training/config/model/transformer.yaml | 5 ++++- 3 files changed, 13 insertions(+), 3 deletions(-) diff --git a/training/src/anemoi/training/config/model/gnn.yaml b/training/src/anemoi/training/config/model/gnn.yaml index 5eda609b..228b6af8 100644 --- 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 index 15a35abf..de558280 100644 --- 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 index 7e36326e..cf5634b3 100644 --- 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