diff --git a/models.py b/models.py
index 8441dea..a15523e 100644
--- a/models.py
+++ b/models.py
@@ -707,6 +707,20 @@ def basic_blocks(dim, index, layers,
 
     return blocks
 
+def window_partition(x, window_size):
+    """
+    Args:
+        x: (B, H, W, C)
+        window_size (int): window size
+
+    Returns:
+        windows: (num_windows*B, window_size, window_size, C)
+    """
+    B, C, H, W = x.shape
+    x = x.view(B, H // window_size, window_size, W // window_size, window_size, C)
+    windows = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, C)
+    return windows
+
 class WindowAttention(nn.Module):
     r""" Window based multi-head self attention (W-MSA) module with relative position bias.
     It supports both of shifted and non-shifted window.
@@ -786,11 +800,6 @@ class WindowAttention(nn.Module):
             x: input features with shape of (num_windows*B, N, C)
             mask: (0/-inf) mask with shape of (num_windows, Wh*Ww, Wh*Ww) or None
         """
-        print(x.shape)
-        B_, C, N, _ = x.shape
-        x = x.reshape(B_, C, N * N)
-        B_, C, N = x.shape
-        x = x.reshape(B_, N, C)
         qkv_bias = None
         if self.q_bias is not None:
             qkv_bias = torch.cat((self.q_bias, torch.zeros_like(self.v_bias, requires_grad=False), self.v_bias))
@@ -843,6 +852,22 @@ class WindowAttention(nn.Module):
         # x = self.proj(x)
         flops += N * self.dim * self.dim
         return flops
+    
+def window_reverse(windows, window_size, H, W):
+    """
+    Args:
+        windows: (num_windows*B, window_size, window_size, C)
+        window_size (int): Window size
+        H (int): Height of image
+        W (int): Width of image
+
+    Returns:
+        x: (B, H, W, C)
+    """
+    B = int(windows.shape[0] / (H * W / window_size / window_size))
+    x = windows.view(B, H // window_size, W // window_size, window_size, window_size, -1)
+    x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, -1, H, W)
+    return x
 
 class PoolFormerBlock(nn.Module):
     """
@@ -868,7 +893,9 @@ 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)
+        self.window_size = 7
+        self.attn_mask = None
+        self.token_mixer = WindowAttention(dim=dim, window_size=to_2tuple(self.window_size), num_heads=10)
         self.norm2 = norm_layer(dim)
         mlp_hidden_dim = int(dim * mlp_ratio)
         self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, 
@@ -885,15 +912,21 @@ class PoolFormerBlock(nn.Module):
                 layer_scale_init_value * torch.ones((dim)), requires_grad=True)
 
     def forward(self, x):
+        B, C, H, W = x.shape
+        x_windows = window_partition(x, self.window_size)
+        x_windows = x_windows.view(-1, self.window_size * self.window_size, C)
+        attn_windows = self.token_mixer(self.norm1(x_windows), mask=self.attn_mask)
+        attn_windows = attn_windows.view(-1, self.window_size, self.window_size, C)
+        x_attn = window_reverse(attn_windows, self.window_size, H, W)
         if self.use_layer_scale:
             x = x + self.drop_path(
                 self.layer_scale_1.unsqueeze(-1).unsqueeze(-1)
-                * self.token_mixer(self.norm1(x)))
+                * x_attn)
             x = x + self.drop_path(
                 self.layer_scale_2.unsqueeze(-1).unsqueeze(-1)
                 * self.mlp(self.norm2(x)))
         else:
-            x = x + self.drop_path(self.token_mixer(self.norm1(x)))
+            x = x + self.drop_path(x_attn)
             x = x + self.drop_path(self.mlp(self.norm2(x)))
         return x
 class PatchEmbed(nn.Module):