def grid_reverse(
        grid: torch.Tensor,
        original_size: Tuple[int, int],
        grid_size: Tuple[int, int] = (7, 7)
) -> torch.Tensor:
    """ Reverses the grid partition.
    Args:
        Grid (torch.Tensor): Grid tensor of the shape [B * grids, grid_size[0], grid_size[1], C].
        original_size (Tuple[int, int]): Original shape.
        grid_size (Tuple[int, int], optional): Grid size which have been applied. Default (7, 7)
    Returns:
        output (torch.Tensor): Folded output tensor of the shape [B, C, original_size[0], original_size[1]].
    """
    # Get height, width, and channels
    (h, w), c = original_size, grid.shape[-1]
    # Compute original batch size
    b = int(grid.shape[0] / (h * w / grid_size[0] / grid_size[1]))
    # Fold grid tensor
    output = grid.view(b, h // grid_size[0], w // grid_size[1], grid_size[0], grid_size[1], c)
    output = output.permute(0, 5, 3, 1, 4, 2).contiguous().view(b, c, h, w)
    return output
Reverses the grid partition.
Args:
Grid (torch.Tensor): Grid tensor of the shape [B * grids, grid_size[0], grid_size[1], C].
original_size (Tuple[int, int]): Original shape.
grid_size (Tuple[int, int], optional): Grid size which have been applied. Default (7, 7)
Returns:
output (torch.Tensor): Folded output tensor of the shape [B, C, original_size[0], original_size[1]].
towhee/models/max_vit/max_vit_utils.py