def img2windows(img, h_split, w_split):
    """Convert input tensor into split stripes

    Args:
        img: tensor, image tensor with shape [B, C, H, W]
        h_split: int, splits width in height direction
        w_split: int, splits width in width direction
    Returns:
        out: tensor, splitted image
    """
    B, C, H, W = img.shape
    out = img.reshape([B, C, H // h_split, h_split, W // w_split, w_split])
    out = out.transpose([0, 2, 4, 3, 5, 1]) # [B, H//h_split, W//w_split, h_split, w_split, C]
    out = out.reshape([-1, h_split * w_split, C]) # [B, H//h_split, W//w_split, h_split*w_split, C]
    return out