TypeError: Invalid shape (416, 528, 1) for image data

[Starfish2018]

Hello, When I try to run the following lines on Colab,

a = DataLoader_Class(path_zip, path)
b = a.read_all_image()
a.see_an_image(5)

I get this error with see_an_image(5):

torch.Size([144, 1, 416, 528])
torch.Size([144, 1, 416, 528])
torch.Size([72, 1, 416, 528])

TypeError Traceback (most recent call last)
in ()
1 a = DataLoader_Class(path_zip, path)
2 b = a.read_all_image()
----> 3 a.see_an_image(5)

5 frames
/usr/local/lib/python3.7/dist-packages/matplotlib/image.py in set_data(self, A)
697 or self._A.ndim == 3 and self._A.shape[-1] in [3, 4]):
698 raise TypeError("Invalid shape {} for image data"
--> 699 .format(self._A.shape))
700
701 if self._A.ndim == 3:

TypeError: Invalid shape (416, 528, 1) for image data

Do you have any idea how I can fix this issue?
Thanks in advance