cleanup region vit

liaowanyi · Nov 22, 2021 · 6665fc6 · 6665fc6
1 parent 5b2382f
commit 6665fc6
Showing 1 changed file with 1 addition and 5 deletions.
diff --git a/vit_pytorch/regionvit.py b/vit_pytorch/regionvit.py
@@ -247,11 +247,7 @@ def __init__(
             nn.Linear(last_dim, num_classes)
         )
 
-    def forward(
-        self,
-        x,
-        return_local_tokens = False
-    ):
+    def forward(self, x):
         *_, h, w = x.shape
         assert divisible_by(h, self.region_patch_size) and divisible_by(w, self.region_patch_size), 'height and width must be divisible by region patch size'
         assert divisible_by(h, self.local_patch_size) and divisible_by(w, self.local_patch_size), 'height and width must be divisible by local patch size'