[Solved] How do you make .make_grid work with a direct 4D tensor input?

In the Pytorch document, it says you can use both a 4D tensor or a List of images as input for make_grid. I tried to use a 4D tensor directly with make_grid with a tensor shape of [10, 1, 64, 64]. So 10 grayscale images that are 64x64. When I call make_grid it would return me the actual values instead of the images. I guess I could then loop thought make_grid and plot the image that way but I am not sure if it is the correct way to use this?

I have 10 images and I would like to plot 2 rows with 5 images in each. How can I do that with a tensor shape of [10, 1, 64, 64]?

My code:

image_pack = image_pack.data.cpu()
print(image_pack.shape) # out (10, 1, 64, 64) 
torchvision.utils.make_grid(image_pack, 5)

Edited: The linked notebook in the document only shows examples if the input is a list of same size images. I am assuming if the input is 4D then it would just automatically work? This is why I am asking this question.

Edited2: Right after I posted this I noticed that I am miss a small helper function imshow which will turn the data into numpy array and switch the channels around and finally plot the images. ops…