Repository: pytorch
Issue: SummaryWriter add_embedding issue with label_img
Labels: oncall: visualization

### 🐛 Describe the bug

Posting working make_sprite function from `utils/tensorboard/_embedding.py`
https://github.com/pytorch/pytorch/blob/master/torch/utils/tensorboard/_embedding.py#L24

```
def make_sprite(label_img, save_path):
    from PIL import Image
    from io import BytesIO

    # this ensures the sprite image has correct dimension as described in
    # https://www.tensorflow.org/get_started/embedding_viz
    print('label_img.size(0):', label_img.size(0))
    nrow = int(math.ceil((label_img.size(0)) ** 0.5))
    print('nrow:', nrow)
    np_imgs = make_np(label_img)
    print('np_imgs:', np_imgs)
    # plt.imshow(np.moveaxis(np_imgs[0], 0, 2))
    # plt.show()
    arranged_img_CHW = make_grid(np_imgs, ncols=nrow)
    print('arranged_img_CHW:', arranged_img_CHW)
    print('arranged_img_CHW.shape:', arranged_img_CHW.shape)
    # plt.imshow(np.moveaxis(arranged_img_CHW, 0, 2))
    # plt.show()

    # augment images so that #images equals nrow*nrow
    arranged_augment_square_HWC = np.zeros((arranged_img_CHW.shape[2], arranged_img_CHW.shape[2], 3))
    print('arranged_augment_square_HWC:', arranged_augment_square_HWC)
    print('arranged_augment_square_HWC.shape:', arranged_augment_square_HWC.shape)
    arranged_img_HWC = arranged_img_CHW.transpose(1, 2, 0)  # chw -> hwc
    print('arranged_img_HWC:', arranged_img_HWC)
    print('arranged_img_HWC.shape:', arranged_img_HWC.shape)
    # plt.imshow(arranged_img_HWC)
    # plt.show()

    arranged_augment_square_HWC[:arranged_img_HWC.shape[0], :, :] = arranged_img_HWC
    print('arranged_augment_square_HWC:', arranged_augment_square_HWC)
    print('arranged_augment_square_HWC.shape:', arranged_augment_square_HWC.shape)

    plt.imshow(arranged_img_HWC)
    plt.show()

    # transformed_img = np.uint8((arranged_augment_square_HWC * 255).clip(0, 255))
    # print(transformed_img)
    # print(transformed_img.shape)

    im = Image.fromarray(arranged_img_HWC)
    # im = Image.fromarray(arranged_augment_square_HWC)

    with BytesIO() as buf:
        im.save(buf, format="PNG")
        im_bytes = buf.getvalue()

    fs = tf.io.gfile.get_filesystem(save_path)
    # print(im_bytes)
    fs.write(fs.join(save_path, 'sprite.png'), im_bytes, binary_mode=True)
```

Commenting this line creates correct sprite and tensorboard is correctly displaying it, seems like it is not required here, please clarify what it should be doing...:
```
arranged_augment_square_HWC[:arranged_img_HWC.shape[0], :, :] = arranged_img_HWC
```
Along with this change, which is not required, because i already have correct rgb integers from 0 to 255:
```
im = Image.fromarray(np.uint8((arranged_augment_square_HWC * 255).clip(0, 255)))
```
Changed to:
```
im = Image.fromarray(arranged_img_HWC)
```

From my perspective, additional asserts are required, to provide useful information when forwarding data into `add_embedding` function, to inform devs about expected data format inside it or create additional `if` statements and check which transformations should be done and which shouldn't in the above two likes of code which i have commented out or changed.


### Versions

Collecting environment information...
PyTorch version: 1.11.0+cu102
Is debug build: False
CUDA used to build PyTorch: 10.2
ROCM used to build PyTorch: N/A

OS: Ubuntu 21.10 (x86_64)
GCC version: (Ubuntu 11.2.0-7ubuntu2) 11.2.0
Clang version: 13.0.0-2
CMake version: version 3.18.4
Libc version: glibc-2.34

Python version: 3.9.7 (default, Sep 16 2021, 13:09:58)  [GCC 7.5.0] (64-bit runtime)
Python platform: Linux-5.17.0-051700-generic-x86_64-with-glibc2.34
Is CUDA available: True
CUDA runtime version: Could not collect
GPU models and configuration: GPU 0: NVIDIA GeForce RTX 3070 Laptop GPU
Nvidia driver version: 510.60.02
cuDNN version: Probably one of the following:
/usr/local/cuda-11.5/targets/x86_64-linux/lib/libcudnn.so.8.3.3
/usr/local/cuda-11.5/targets/x86_64-linux/lib/libcudnn_adv_infer.so.8.3.3
/usr/local/cuda-11.5/targets/x86_64-linux/lib/libcudnn_adv_train.so.8.3.3
/usr/local/cuda-11.5/targets/x86_64-linux/lib/libcudnn_cnn_infer.so.8.3.3
/usr/local/cuda-11.5/targets/x86_64-linux/lib/libcudnn_cnn_train.so.8.3.3
/usr/local/cuda-11.5/targets/x86_64-linux/lib/libcudnn_ops_infer.so.8.3.3
/usr/local/cuda-11.5/targets/x86_64-linux/lib/libcudnn_ops_train.so.8.3.3
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] facenet-pytorch==2.5.2
[pip3] mypy-extensions==0.4.3
[pip3] numpy==1.23.0
[pip3] numpydoc==1.1.0
[pip3] torch==1.11.0
[pip3] torchaudio==0.11.0
[pip3] torchvision==0.12.0
[conda] blas                      1.0                         mkl  
[conda] cudatoolkit               11.3.1               h2bc3f7f_2  
[conda] facenet-pytorch           2.5.2                    pypi_0    pypi
[conda] ffmpeg                    4.3                  hf484d3e_0    pytorch
[conda] mkl                       2021.4.0           h06a4308_640  
[conda] mkl-service               2.4.0            py39h7f8727e_0  
[conda] mkl_fft                   1.3.1            py39hd3c417c_0  
[conda] mkl_random                1.2.2            py39h51133e4_0  
[conda] numpy                     1.23.0                   pypi_0    pypi
[conda] numpydoc                  1.1.0              pyhd3eb1b0_1  
[conda] pytorch-mutex             1.0                        cuda    pytorch
[conda] torch                     1.11.0                   pypi_0    pypi
[conda] torchaudio                0.11.0               py39_cu113    pytorch
[conda] torchvision               0.12.0               py39_cu113    pytorch