diff --git a/models.py b/models.py
index 73c5f3b..7aa5352 100644
--- a/models.py
+++ b/models.py
@@ -866,7 +866,7 @@ class PoolFormerBlock(nn.Module):
         self.norm1 = norm_layer(dim)
         #self.token_mixer = Pooling(pool_size=pool_size)
         # self.token_mixer = FNetBlock()
-        self.token_mixer = WindowAttention(dim=dim, window_size=to_2tuple(7), num_heads=10, pretrained_window_size=[5,5])
+        self.token_mixer = WindowAttention(dim=dim, window_size=to_2tuple(7), num_heads=3)
         self.norm2 = norm_layer(dim)
         mlp_hidden_dim = int(dim * mlp_ratio)
         self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim,