fix: when window_size is passes as array (#26800)

huggingface · Oct 17, 2023 · 85e9d64 · 85e9d64
1 parent b3961f7
commit 85e9d64
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/src/transformers/models/swinv2/modeling_swinv2.py b/src/transformers/models/swinv2/modeling_swinv2.py
@@ -791,14 +791,19 @@ def __init__(
         super().__init__()
         self.config = config
         self.dim = dim
+        window_size = (
+            config.window_size
+            if isinstance(config.window_size, collections.abc.Iterable)
+            else (config.window_size, config.window_size)
+        )
         self.blocks = nn.ModuleList(
             [
                 Swinv2Layer(
                     config=config,
                     dim=dim,
                     input_resolution=input_resolution,
                     num_heads=num_heads,
-                    shift_size=0 if (i % 2 == 0) else config.window_size // 2,
+                    shift_size=[0, 0] if (i % 2 == 0) else [window_size[0] // 2, window_size[1] // 2],
                     pretrained_window_size=pretrained_window_size,
                 )
                 for i in range(depth)