Assertion Error in : images.ndim==4 is not true for the given dataset in github.

[priyamp2309]

Created a dataset of size = 1 with image resolution (192, 320)
Using the training regime with segmentation masks
Created Generator with 5887448 parameters
Created Discriminator (3+4 blocks) with 2443581 parameters
Traceback (most recent call last):
File "train.py", line 33, in
out_G = diff_augment(out_G)
File "/home/usd.local/priyam.pandey/priyam/code/one-shot-synthesis-master/one-shot-synthesis-master/core/differentiable_augmentation/diff_augm.py", line 12, in call
return self.augment_func(batch)
File "/home/usd.local/priyam.pandey/anaconda3/envs/osmis/lib/python3.8/site-packages/torch/nn/modules/module.py", line 722, in _call_impl
result = self.forward(*input, **kwargs)
File "/home/usd.local/priyam.pandey/priyam/code/one-shot-synthesis-master/one-shot-synthesis-master/core/differentiable_augmentation/augment.py", line 182, in forward
assert isinstance(images, torch.Tensor) and images.ndim == 4
AssertionError

[pietrobert3]

I have the same problem