Как вывести изображение в pytorch из обесцвеченного тензора?

Пытаюсь загрузить свои данные для обучения. Делаю таким образом.

batch_sizeDif = len(imagDif_Tenzor)
batchDif = torch.zeros(batch_sizeDif, 3,  493, 733, dtype=torch.uint8)
    
batch_sizeD = len(imagD_Tenzor)
batchD = torch.zeros(batch_sizeD, 3,  493, 733, dtype=torch.uint8)

for i, filename in enumerate(imagDif_Tenzor):
 batchDif[i] = torchvision.io.read_image(os.path.join(imagDif_path, filename))

batchDif = setup_aug_tfms([Saturation(max_lighting=0.0, p=1.0, draw=0.0)])

for i, filename in enumerate(imagD_Tenzor):
 batchD[i] = torchvision.io.read_image(os.path.join(imagD_path, filename))

Прошу заметить третью строку с конца. Я обесцвечиваю тензор. И после перестаёт работать следующий код.

fig = plt.figure()
ax1 = fig.add_subplot(2,2,1)
ax1.imshow(batchDif[0].permute(1, 2, 0))
ax2 = fig.add_subplot(2,2,2)
ax2.imshow(batchD[0].permute(1, 2, 0))
ax3 = fig.add_subplot(2,2,3)
ax3.imshow(batchDif[20].permute(1, 2, 0))
ax4 = fig.add_subplot(2,2,4)
ax4.imshow(batchD[20].permute(1, 2, 0))
plt.show()

тут ошибка в 3 строке с верху. 'Saturation' object has no attribute 'permute'.

Какой атрибут использовать чтобы вывести изображение? Или вообще по другому действовать? Если не обесцвечивать то изображения выводятся без ошибок.