🐛 add supported attention

jgrss · Jul 31, 2024 · 85dc5d0 · 85dc5d0
1 parent c685fa5
commit 85dc5d0
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/src/cultionet/nn/modules/convolution.py b/src/cultionet/nn/modules/convolution.py
@@ -356,6 +356,7 @@ def __init__(
             assert self.attention_weights in [
                 AttentionTypes.FRACTAL,
                 AttentionTypes.SPATIAL_CHANNEL,
+                AttentionTypes.NATTEN,
             ], "The attention method is not supported."
 
             self.gamma = nn.Parameter(torch.ones(1))