correct shape (BS,H,W,C) not working in torchvision.utils.save_image

Let’s BS be the batch size, H the height, w the weight, and c the number of channels which is 3 in my case.

When I save my image in this shape (BS,C,H,W) with

torchvision.utils.save_image(image, path)

it works very well but the image is unreadable since the format is wrong.

But when I am reshaping my image into the right format which is (BS,H,W,C), The code below does not work

image = image.reshape(BS,H,W,C)  
torchvision.utils.save_image(image, path)

Here is the error that I am stuck in:

TypeError: Cannot handle this data type: (1, 1, 256), |u1

Answer

You wanted your image to have size (BS, C, H, W), but you are incorrectly reshaping it.

Assuming the image.shape is (BS, H, W, C), perhaps you meant to perform image = image.permute(0, 3, 1, 2) which moves the channel dimension in the second position, to obtain the required shape.