try modify swin

2024-04-29 20:48:47 +07:00
parent d3a6cfe041
commit a14267d96c
1 changed files with 1 additions and 1 deletions
--- a/models.py
+++ b/models.py
@@ -535,7 +535,7 @@ class FouriER(torch.nn.Module):
        layers = [2, 2, 6, 2]
        embed_dims = [self.p.embed_dim, 320, 256, 128]
        mlp_ratios = [4, 4, 8, 12]
-        num_heads = [2, 4, 8, 16]
+        num_heads = [4, 4, 4, 4]
        downsamples = [True, True, True, True]
        pool_size=3
        act_layer=nn.GELU