class_docstr: Partition the input tensor into non-overlapping windows.
docstr: Args:
      x (Tensor): Input tensor with expected layout of [B, C, H, W].
      p (int): Number of partitions.
  Returns:
      Tensor: Output tensor with expected layout of [B, H/P, W/P, P*P, C].
docstr_tok: ["Args",":","x","(","Tensor",")",":","Input","tensor","with","expected","layout","of","[","B",",","C",",","H",",","W","]",".","p","(","int",")",":","Number","of","partitions",".","Returns",":","Tensor",":","Output","tensor","with","expected","layout","of","[","B",",","H\/P",",","W\/P",",","P","*","P",",","C","]","."]
code: def forward(self, x: Tensor, p: int) -> Tensor:
      """
      Args:
          x (Tensor): Input tensor with expected layout of [B, C, H, W].
          p (int): Number of partitions.
      Returns:
          Tensor: Output tensor with expected layout of [B, H/P, W/P, P*P, C].
      """
      B, C, H, W = x.shape
      P = p
      # chunk up H and W dimensions
      x = x.reshape(B, C, H // P, P, W // P, P)
      x = x.permute(0, 2, 4, 3, 5, 1)
      # colapse P * P dimension
      x = x.reshape(B, (H // P) * (W // P), P * P, C)
      return x
  
imports: import math
  from collections import OrderedDict
  from functools import partial
  from typing import Any, Callable, List, Optional, Sequence, Tuple
  import numpy
  import torch
  import torch.nn.functional
  from torch import nn, Tensor
  from torchvision.models._api import register_model, Weights, WeightsEnum
  from torchvision.models._meta import _IMAGENET_CATEGORIES
  from torchvision.models._utils import _ovewrite_named_param, handle_legacy_interface
  from torchvision.ops.misc import Conv2dNormActivation, SqueezeExcitation
  from torchvision.ops.stochastic_depth import StochasticDepth
  from torchvision.transforms._presets import ImageClassification, InterpolationMode
  from torchvision.utils import _log_api_usage_once
label_desc: Use image node_id 2 for calling the WindowPartition obj's underlying member method code with example usage: obj.forward(x, p) and returns: x