Hello, I might be a bit late, but you may use torchvision.utils.make_grid() to accomplish this.
Example code:
class TensorBoardHelper:

    def __init__(self, name=None):
        self.writer = SummaryWriter(name)


    def add_image_grid(self, images, tag, global_step):
        images = images.detach()
        img_grid = torchvision.utils.make_grid(images, 3, normalize=True, scale_each=True,
                padding=2, pad_value=1.0, # Pad 2 pixels, with white borders
                )
        self.writer.add_image(tag, img_grid, global_step=global_step)

Output:
image403×528 17.9 KB
https://pytorch.org/docs/stable/torchvision/utils.html#torchvision.utils.make_grid