io.write_png saves RGB image when given greyscale tensor

🐛 Bug

When giving a 3D 1xHxW uint8 tensor to torchvision.io.write_png, it saves it as a coloured PNG instead of a greyscale PNG. torchvision.io.write_jpeg works correctly.

To Reproduce

Steps to reproduce the behavior:

1. Create a 3D 1x64x64 uint8 tensor.
2. Give it to torchvision.io.write_png.
3. Check the saved PNG’s attributes, e.g. using identify from ImageMagick. It should say sRGB instead of greyscale.

import torch
from torchvision.io import write_png

img = (torch.rand(1, 64, 64) * 255).byte()
write_png(img, "test.png")

Expected behavior

The saved PNG should be greyscale.

Environment

Please copy and paste the output from our environment collection script (or fill out the checklist below manually).

You can get the script and run it with:

wget https://raw.githubusercontent.com/pytorch/pytorch/master/torch/utils/collect_env.py
# For security purposes, please check the contents of collect_env.py before running it.
python collect_env.py

PyTorch version: 1.8.0
Is debug build: False
CUDA used to build PyTorch: 10.2
ROCM used to build PyTorch: N/A

OS: Arch Linux (x86_64)
GCC version: (GCC) 10.2.0
Clang version: 11.1.0
CMake version: Could not collect

Python version: 3.9 (64-bit runtime)
Is CUDA available: True
CUDA runtime version: 11.2.142
GPU models and configuration: GPU 0: GeForce GTX 960M
Nvidia driver version: 460.56
cuDNN version: Probably one of the following:
/usr/lib/libcudnn.so.8.1.0
/usr/lib/libcudnn_adv_infer.so.8.1.0
/usr/lib/libcudnn_adv_train.so.8.1.0
/usr/lib/libcudnn_cnn_infer.so.8.1.0
/usr/lib/libcudnn_cnn_train.so.8.1.0
/usr/lib/libcudnn_ops_infer.so.8.1.0
/usr/lib/libcudnn_ops_train.so.8.1.0
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] numpy==1.19.5
[pip3] torch==1.8.0
[pip3] torchvision==0.9.0
[conda] Could not collect

[fmassa]

@rharish101 oh you are right, the ternary operator was not implemented correctly there and changing this will probably fix the issue

[fmassa]

@andfoy could you have a look?