Repository: pytorch/vision

Issue:
AutoAugment's document is wrong (or unexpected?)
### 🐛 Describe the bug

In the document, it is said that:
> [...] If the image is torch Tensor, it should be of type torch.uint8, and it is expected to have […, 1 or 3, H, W] shape, where … means an arbitrary number of leading dimensions.

So, if I understand correctly, the shape `(1,H,W)` still works?

However, when I pass an input tensor whose shape is `(1,H,W)`, it sometimes does the job and sometimes doesn't.

To reproduce:

First, since it is hard to reproduce a randomized transformation that I have just suddenly met, I have to loop through several random seeds using the code below:

```python
for seed in range(0, 1000):
    print(seed)
    torch.manual_seed(seed)
    image = torch.randint(0,255,(1,32,64), dtype=torch.uint8)
    tf = transforms.AutoAugment(transforms.autoaugment.AutoAugmentPolicy.CIFAR10)
    image = tf(image)
    image
```

And the error raises when `seed=5`, I then hard-code the seed as follows:
```python
torch.manual_seed(5)
image = torch.randint(0,255,(1,32,64), dtype=torch.uint8)
tf = transforms.AutoAugment(transforms.autoaugment.AutoAugmentPolicy.CIFAR10)
image = tf(image)
image
```

The traceback error is:
```
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
/tmp/ipykernel_32333/3324237737.py in <module>
      2 image = torch.randint(0,255,(1,1,32,64), dtype=torch.uint8)
      3 tf = transforms.AutoAugment(transforms.autoaugment.AutoAugmentPolicy.CIFAR10)
----> 4 image = tf(image)
      5 image

~/miniconda3/envs/cv/lib/python3.8/site-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
   1049         if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
   1050                 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1051             return forward_call(*input, **kwargs)
   1052         # Do not call functions when jit is used
   1053         full_backward_hooks, non_full_backward_hooks = [], []

~/miniconda3/envs/cv/lib/python3.8/site-packages/torchvision/transforms/autoaugment.py in forward(self, img)
    214                     img = F.adjust_brightness(img, 1.0 + magnitude)
    215                 elif op_name == "Color":
--> 216                     img = F.adjust_saturation(img, 1.0 + magnitude)
    217                 elif op_name == "Contrast":
    218                     img = F.adjust_contrast(img, 1.0 + magnitude)

~/miniconda3/envs/cv/lib/python3.8/site-packages/torchvision/transforms/functional.py in adjust_saturation(img, saturation_factor)
    810         return F_pil.adjust_saturation(img, saturation_factor)
    811 
--> 812     return F_t.adjust_saturation(img, saturation_factor)
    813 
    814 

~/miniconda3/envs/cv/lib/python3.8/site-packages/torchvision/transforms/functional_tensor.py in adjust_saturation(img, saturation_factor)
    213     _assert_image_tensor(img)
    214 
--> 215     _assert_channels(img, [3])
    216 
    217     return _blend(img, rgb_to_grayscale(img), saturation_factor)

~/miniconda3/envs/cv/lib/python3.8/site-packages/torchvision/transforms/functional_tensor.py in _assert_channels(img, permitted)
     53     c = _get_image_num_channels(img)
     54     if c not in permitted:
---> 55         raise TypeError("Input image tensor permitted channel values are {}, but found {}".format(permitted, c))
     56 
     57 

TypeError: Input image tensor permitted channel values are [3], but found 1
```

I think some operators do not work (or it hasn't been designed to work) with grayscale images?
I still randomly meet this error with IMAGENET policy also.

### Versions

```
Collecting environment information...
PyTorch version: 1.9.0
Is debug build: False
CUDA used to build PyTorch: 10.2
ROCM used to build PyTorch: N/A

OS: Ubuntu 18.04.5 LTS (x86_64)
GCC version: (Ubuntu 7.5.0-3ubuntu1~18.04) 7.5.0
Clang version: Could not collect
CMake version: Could not collect
Libc version: glibc-2.27

Python version: 3.8.11 (default, Aug  3 2021, 15:09:35)  [GCC 7.5.0] (64-bit runtime)
Python platform: Linux-5.0.0-23-generic-x86_64-with-glibc2.17
Is CUDA available: True
CUDA runtime version: Could not collect
GPU models and configuration: GPU 0: NVIDIA GeForce RTX 2060
Nvidia driver version: 470.57.02
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] numpy==1.20.3
[pip3] torch==1.9.0
[pip3] torchvision==0.10.0
[conda] blas                      1.0                         mkl  
[conda] cudatoolkit               10.2.89              hfd86e86_1  
[conda] ffmpeg                    4.3                  hf484d3e_0    pytorch
[conda] mkl                       2021.3.0           h06a4308_520  
[conda] mkl-service               2.4.0            py38h497a2fe_0    conda-forge
[conda] mkl_fft                   1.3.0            py38h42c9631_2  
[conda] mkl_random                1.2.2            py38h1abd341_0    conda-forge
[conda] numpy                     1.20.3           py38hf144106_0  
[conda] numpy-base                1.20.3           py38h74d4b33_0  
[conda] pytorch                   1.9.0           py3.8_cuda10.2_cudnn7.6.5_0    pytorch
[conda] torchvision               0.10.0               py38_cu102    pytorch
```

cc @vfdev-5 @datumbox


Discussion:
@VinhLoiIT You are absolutely right, the documentation is wrong. Only RGB images are supported.

The reason this works sometimes is due to the randomness of AA, which leads to choosing transforms that can be applied on non RGB images. But then when a transformation like `adjust_saturation` is sampled that requires RGB, it fails.

The doc string is wrong on all 3 autoaugment methods:
https://github.com/pytorch/vision/blob/d9e6d60f49e650c2f70b2cf457b2aca081239930/torchvision/transforms/autoaugment.py#L68

https://github.com/pytorch/vision/blob/d9e6d60f49e650c2f70b2cf457b2aca081239930/torchvision/transforms/autoaugment.py#L252

https://github.com/pytorch/vision/blob/d9e6d60f49e650c2f70b2cf457b2aca081239930/torchvision/transforms/autoaugment.py#L336

We should be saying that only `[..., 3, H, W]` shape is supported. Would you be interested in sending a PR that fixes the problem?


Resolving patch:
diff --git a/torchvision/transforms/autoaugment.py b/torchvision/transforms/autoaugment.py
--- a/torchvision/transforms/autoaugment.py
+++ b/torchvision/transforms/autoaugment.py
@@ -65,8 +65,8 @@ class AutoAugment(torch.nn.Module):
     r"""AutoAugment data augmentation method based on
     `"AutoAugment: Learning Augmentation Strategies from Data" <https://arxiv.org/pdf/1805.09501.pdf>`_.
     If the image is torch Tensor, it should be of type torch.uint8, and it is expected
-    to have [..., 1 or 3, H, W] shape, where ... means an arbitrary number of leading dimensions.
-    If img is PIL Image, it is expected to be in mode "L" or "RGB".
+    to have [..., 3, H, W] shape, where ... means an arbitrary number of leading dimensions.
+    If img is PIL Image, it is expected to be in mode "RGB".
 
     Args:
         policy (AutoAugmentPolicy): Desired policy enum defined by
@@ -249,8 +249,8 @@ class RandAugment(torch.nn.Module):
     `"RandAugment: Practical automated data augmentation with a reduced search space"
     <https://arxiv.org/abs/1909.13719>`_.
     If the image is torch Tensor, it should be of type torch.uint8, and it is expected
-    to have [..., 1 or 3, H, W] shape, where ... means an arbitrary number of leading dimensions.
-    If img is PIL Image, it is expected to be in mode "L" or "RGB".
+    to have [..., 3, H, W] shape, where ... means an arbitrary number of leading dimensions.
+    If img is PIL Image, it is expected to be in mode "RGB".
 
     Args:
         num_ops (int): Number of augmentation transformations to apply sequentially.
@@ -333,8 +333,8 @@ class TrivialAugmentWide(torch.nn.Module):
     r"""Dataset-independent data-augmentation with TrivialAugment Wide, as described in
     `"TrivialAugment: Tuning-free Yet State-of-the-Art Data Augmentation" <https://arxiv.org/abs/2103.10158>`.
     If the image is torch Tensor, it should be of type torch.uint8, and it is expected
-    to have [..., 1 or 3, H, W] shape, where ... means an arbitrary number of leading dimensions.
-    If img is PIL Image, it is expected to be in mode "L" or "RGB".
+    to have [..., 3, H, W] shape, where ... means an arbitrary number of leading dimensions.
+    If img is PIL Image, it is expected to be in mode "RGB".
 
     Args:
         num_magnitude_bins (int): The number of different magnitude values.


Test patch: