code: def grid_partition(
      data: torch.Tensor, grid_size: Tuple[int, int] = (7, 7)
  ) -> torch.Tensor:
      """Grid partition function.
      Args:
          data (torch.Tensor): Input tensor of the shape [B, C, H, W].
          grid_size (Tuple[int, int], optional): Grid size to be applied. Default (7, 7)
      Returns:
          grid (torch.Tensor): Unfolded input tensor of the shape [B * grids, grid_size[0], grid_size[1], C].
      """
      # Get size of input
      b, c, h, w = data.shape
      # Unfold input
      grid = data.view(
          b,
          c,
          grid_size[0],
          h // grid_size[0],
          grid_size[1],
          w // grid_size[1],
      )
      # Permute and reshape [B * (H // grid_size[0]) * (W // grid_size[1]), grid_size[0], window_size[1], C]
      grid = (
          grid.permute(0, 3, 5, 2, 4, 1)
          .contiguous()
          .view(-1, grid_size[0], grid_size[1], c)
      )
      return grid
  
imports: from typing import Tuple
  import torch
label_desc: Use image node_id 3 for calling a global function with example usage: grid_partition(data, grid_size) and returns: grid