class CifarDataset(Dataset):
    """
    :root_indexes:  path to indexes compressed file
    :root_data:     path to data compressed file
    :transforms_:   list of functions from torchvision.transforms
    """

    def __init__(self, root_indexes: str, root_data: str, transforms_: list = []):
        self.transform = transforms.Compose(transforms_)

        r_data = np.load(root_data)

        self.indexes = np.load(root_indexes)['all_indexes']
        self.labels = r_data['all_labels']

        # the line bellow solved my problem
        self.data = np.transpose(r_data['all_data'].reshape(-1,3,32,32), (0,2,3,1))

        self.data = torch.stack([self.transform(im) for im in self.data])

    def __getitem__(self, index) -> tuple:
        idx = self.indexes[index]
        return self.data[idx], self.labels[idx]

    def __len__(self) -> int:
        return len(self.indexes)

It seems the problem was a bit far away from the the grid as I thought so.