import torch


def divide_img(img, patch_size):
    """Divides image into tensor of image patches
    Args:
        img: batch of images, torch.tensor, e.g. [batch_size, channels, 32, 32]
        patch_size: patch size, tuple e.g. (4,4)
    
    Returns:
        A torch.tensor of stacked stacked flattened image patches, e.g. [batch_size, 64, channels, 4, 4]
    """
    height = img.shape[2] 
    width = img.shape[3]
    patch_size_h = patch_size[0]
    patch_size_w = patch_size[1]
    A = []
    for i in range(int(height/patch_size_h)):
        for j in range(int(width/patch_size_w)):
            A.append(img[:,:,i*patch_size_h:i*patch_size_h+patch_size_h,j*patch_size_w:j*patch_size_w+patch_size_w])
    return torch.stack(A).permute(1,0,2,3,4)