diff --git a/models.py b/models.py index 272f53a..3302431 100644 --- a/models.py +++ b/models.py @@ -488,7 +488,7 @@ class FouriER(torch.nn.Module): self.patch_embed = PatchEmbed(in_chans=channels, patch_size=self.p.patch_size, embed_dim=self.p.embed_dim, stride=4, padding=2) network = [] - layers = [4, 4, 12, 4] + 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]