diff --git a/timm/models/swin_transformer.py b/timm/models/swin_transformer.py
index ac8f9d7d..bd396557 100644
--- a/timm/models/swin_transformer.py
+++ b/timm/models/swin_transformer.py
@@ -41,7 +41,7 @@ _int_or_tuple_2_t = Union[int, Tuple[int, int]]
 def window_partition(
         x: torch.Tensor,
         window_size: Tuple[int, int],
-) -> Tuple[torch.Tensor, Tuple[int, int]]:
+) -> torch.Tensor:
     """
     Partition into non-overlapping windows with padding if needed.
     Args: