Add pretrained window size value from config in SwinV2Encoder class

huggingface · Jun 16, 2022 · 66d4513 · 66d4513
1 parent 2547f45
commit 66d4513
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/src/transformers/models/swinv2/modeling_swinv2.py b/src/transformers/models/swinv2/modeling_swinv2.py
@@ -800,6 +800,8 @@ def __init__(self, config, grid_size, pretrained_window_sizes=(0, 0, 0, 0)):
         super().__init__()
         self.num_layers = len(config.depths)
         self.config = config
+        if self.config.pretrained_window_sizes is not None:
+            pretrained_window_sizes = config.pretrained_window_sizes
         dpr = [x.item() for x in torch.linspace(0, config.drop_path_rate, sum(config.depths))]
         self.layers = nn.ModuleList(
             [