using ImageFolder with albumentations in pytorch

I have a situation where I need to use ImageFolder with the albumentations lib to make the augmentations in pytorch - custom dataloader is not an option.

To this end, I am stumped and I am not able to get ImageFolder to work with albumenations. I have tried something along these lines:

class Transforms:
    def __init__(self, transforms: A.Compose):
        self.transforms = transforms

    def __call__(self, img, *args, **kwargs):
        return self.transforms(image=np.array(img))['image']

and then:

    trainset = datasets.ImageFolder(traindir,transform=Transforms(transforms=A.Resize(32 , 32)))

where traindir is some dir with images. I however get thrown a weird error:

RuntimeError: Given groups=1, weight of size [16, 3, 3, 3], expected input[1024, 32, 32, 3] to have 3 channels, but got 32 channels instead

and I cant seem to find a reproducible example to make a simple aug pipleline work with imagefolder.

UPDATE On the recommendation of @Shai, I have done this now:

class Transforms:
    def __init__(self):
        self.transforms = A.Compose([A.Resize(224,224),ToTensorV2()])

    def __call__(self, img, *args, **kwargs):
        return self.transforms(image=np.array(img))['image']

trainset = datasets.ImageFolder(traindir,transform=Transforms())

but I get thrown:

    self.padding, self.dilation, self.groups)
RuntimeError: Input type (torch.cuda.ByteTensor) and weight type (torch.cuda.FloatTensor) should be the same

You need to use ToTensorV2 transformation as the final one:

trainset = datasets.ImageFolder(traindir,transform=Transforms(transforms=A.Compose([A.Resize(32 , 32), ToTensorV2()]))