This commit is contained in:
thanhvc3 2024-04-27 11:51:35 +07:00
parent dda7f13dbd
commit ae0f43ab4d

View File

@ -894,7 +894,7 @@ class PoolFormerBlock(nn.Module):
self.norm1 = norm_layer(dim) self.norm1 = norm_layer(dim)
#self.token_mixer = Pooling(pool_size=pool_size) #self.token_mixer = Pooling(pool_size=pool_size)
# self.token_mixer = FNetBlock() # self.token_mixer = FNetBlock()
self.window_size = 8 self.window_size = 4
self.attn_mask = None self.attn_mask = None
self.token_mixer = WindowAttention(dim=dim, window_size=to_2tuple(self.window_size), num_heads=4) self.token_mixer = WindowAttention(dim=dim, window_size=to_2tuple(self.window_size), num_heads=4)
self.norm2 = norm_layer(dim) self.norm2 = norm_layer(dim)