from PIL import Image
// this ensures the sprite image has correct dimension as described in
// https://www.tensorflow.org/get_started/embedding_viz
nrow = int(math.ceil((label_img.size(0)) ** 0.5))
arranged_img_CHW = make_grid(make_np(label_img), ncols=nrow)
// augment images so that //images equals nrow*nrow
arranged_augment_square_HWC = np.ndarray((arranged_img_CHW.shape[2], arranged_img_CHW.shape[2], 3))
After Change
// as images to be shown.
// args: label_img: tensor in NCHW
assert label_img.shape[2] == label_img.shape[3], "Image should be square, see tensorflow/tensorboard/먶"
total_pixels = label_img.shape[0] * label_img.shape[2] * label_img.shape[3]
pixels_one_side = total_pixels ** 0.5
number_of_images_per_row = int(math.ceil(pixels_one_side / label_img.shape[3]))
arranged_img_CHW = make_grid(make_np(label_img), ncols=number_of_images_per_row)