Creating an image grid of tensors out of a batch using matplotlib/pytorch-CodePudding

I am trying to create a grid of images (e.g. 3 by 3) from a batch of tensors that will be fed into a GAN through a data loader in the next step. With the below code I was able to transform the tensors into images that are displayed in a grid in the right position. The problem is, that they are all displayed in a separate grid as shown here:

Please note that the code allows you to generate all the images you want, not only 3 times 3. I have a folder called cryptopunks with a lot of images called #.png (e.g., 1.png, ..., 34.png, ...). Just change the row_count and col_count variable values. For instance, for row_count=6 and col_count=8 you get:

If your image files do not have that naming pattern above (i.e., just random names), just replace the first lines with the following ones:

import os
from pathlib import Path

import matplotlib.pyplot as plt
from mpl_toolkits.axes_grid1 import ImageGrid
import matplotlib.image as mpimg

for root, _, filenames in os.walk("cryptopunks/"):
    cryptopunks = [
        mpimg.imread(Path(root, filename)) for filename in filenames
    ]

row_count = 3
col_count = 3

# Here same code as above.

(I have downloaded the CryptoPunks dataset from