From b9273b6696337d5754b006ee640f8082b01fe479 Mon Sep 17 00:00:00 2001
From: thanhvc3 <thanhvc3@vng.com.vn>
Date: Sun, 28 Apr 2024 15:27:41 +0700
Subject: [PATCH] try gtp vit

---
 models.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/models.py b/models.py
index 598fb5c..2bdc97b 100644
--- a/models.py
+++ b/models.py
@@ -1238,7 +1238,7 @@ class PoolFormerBlock(nn.Module):
             self.layer_scale_2 = nn.Parameter(
                 layer_scale_init_value * torch.ones((dim)), requires_grad=True)
 
-    def forward(self, x, weight):
+    def forward(self, x, weight, token_scales = None):
         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)