Repository: pytorch
Issue: Some operations do not keep `channels_last` memory format which yields accuracy drop
Labels: triaged,module: memory format

### 🐛 Describe the bug

I tried to switch to `torch.channels_last` in order to reduce the training time but I noticed that the new training shows quite a performance drop. While looking deeper into the cause, I noticed that there are some modules which do not keep the `torch.channels_last` memory format:

```python
from functools import partial
import itertools
import torch
import torch.nn as nn

devices = ['cpu', 'cuda']
modules = [partial(nn.ConvTranspose2d, in_channels=3, out_channels=3, kernel_size=(2, 2), stride=(2, 2), bias=False),
           partial(nn.Conv2d, in_channels=3, out_channels=3, kernel_size=(2, 2), stride=(2, 2), bias=False),
           partial(nn.Linear, in_features=5, out_features=5, bias=False),
           partial(nn.Conv2d, in_channels=3, out_channels=3, kernel_size=(3, 3), stride=1, padding=(1, 1), bias=False),
           partial(nn.Conv2d, in_channels=3, out_channels=3, kernel_size=(1, 1)),
           partial(nn.Conv2d, in_channels=3, out_channels=3, kernel_size=(1, 2), stride=1, bias=False),
           partial(nn.Conv2d, in_channels=3, out_channels=3, kernel_size=(1, 2), stride=1, padding=(0, 1)),
           nn.GELU,
           partial(nn.BatchNorm2d, num_features=3)]

for device, module in itertools.product(devices, modules):
    m = module()

    # create random input
    inp = torch.randn(2, 3, 4, 5)

    # cast input and module to device/memory format
    inp = inp.to(memory_format=torch.channels_last).to(device=device)
    m = m.to(memory_format=torch.channels_last).to(device=device)

    out = m(inp)

    if not out.is_contiguous(memory_format=torch.channels_last):
        print(f'Operation does not keep channels_last for '
              f'device = {device}, module = {module} but has now stride {out.stride()}')
```

which yields

```bash
Operation does not keep channels_last for device = cpu, module = functools.partial(<class 'torch.nn.modules.conv.ConvTranspose2d'>, in_channels=3, out_channels=3, kernel_size=(2, 2), stride=(2, 2), bias=False) but has now stride (240, 80, 10, 1)
Operation does not keep channels_last for device = cpu, module = functools.partial(<class 'torch.nn.modules.linear.Linear'>, in_features=5, out_features=5, bias=False) but has now stride (60, 20, 5, 1)
Operation does not keep channels_last for device = cuda, module = functools.partial(<class 'torch.nn.modules.linear.Linear'>, in_features=5, out_features=5, bias=False) but has now stride (60, 20, 5, 1)

```

As `nn.Linear` does not keep the channels_last memory layout on the GPU, this then causes a problem, as the following modules are now in `torch.channels_last` but the input is not anymore:

```python
from functools import partial
import itertools
import torch
import torch.nn as nn

devices = ['cpu', 'cuda']
memory_formats = [torch.contiguous_format, torch.channels_last]
modules = [partial(nn.ConvTranspose2d, in_channels=3, out_channels=3, kernel_size=(2, 2), stride=(2, 2), bias=False),
           partial(nn.Conv2d, in_channels=3, out_channels=3, kernel_size=(2, 2), stride=(2, 2), bias=False),
           partial(nn.Linear, in_features=5, out_features=5, bias=False),
           partial(nn.Conv2d, in_channels=3, out_channels=3, kernel_size=(3, 3), stride=1, padding=(1, 1), bias=False),
           partial(nn.Conv2d, in_channels=3, out_channels=3, kernel_size=(1, 1)),
           partial(nn.Conv2d, in_channels=3, out_channels=3, kernel_size=(1, 2), stride=1, bias=False),
           partial(nn.Conv2d, in_channels=3, out_channels=3, kernel_size=(1, 2), stride=1, padding=(0, 1)),
           nn.GELU,
           partial(nn.BatchNorm2d, num_features=3)]

for device, memory_format, module in itertools.product(devices, memory_formats, modules):
    m1 = module()
    m2 = module()

    # use identical weights for both modules
    m2.load_state_dict(m1.state_dict())

    # create random input
    inp = torch.randn(2, 3, 4, 5)

    # cast module to device/memory format
    m1 = m1.to(device=device)
    m2 = m2.to(memory_format=memory_format).to(device=device)

    # cast input to device (keeping `contiguous` memory format)
    inp = inp.to(device=device)

    out1 = m1(inp)
    out2 = m2(inp)
    if not torch.allclose(out1, out2):
        print(f'Results differ for device = {device}, memory_format = {memory_format}, module = {module} '
              f'with max difference {torch.max(torch.abs(out1 - out2)).item()}')
```

which is a problem for the `Conv2d`:

```bash
Results differ for device = cuda, memory_format = torch.channels_last, module = functools.partial(<class 'torch.nn.modules.conv.Conv2d'>, in_channels=3, out_channels=3, kernel_size=(2, 2), stride=(2, 2), bias=False) with max difference 0.0002925097942352295
Results differ for device = cuda, memory_format = torch.channels_last, module = functools.partial(<class 'torch.nn.modules.conv.Conv2d'>, in_channels=3, out_channels=3, kernel_size=(1, 1)) with max difference 0.0007414817810058594
Results differ for device = cuda, memory_format = torch.channels_last, module = functools.partial(<class 'torch.nn.modules.conv.Conv2d'>, in_channels=3, out_channels=3, kernel_size=(1, 2), stride=1, bias=False) with max difference 0.000747382640838623
Results differ for device = cuda, memory_format = torch.channels_last, module = functools.partial(<class 'torch.nn.modules.conv.Conv2d'>, in_channels=3, out_channels=3, kernel_size=(1, 2), stride=1, padding=(0, 1)) with max difference 0.0005008578300476074
```

So I think the `nn.Conv2d` silently assumes that the input has the correct memory layout which is not the case due to a previous module that brought it back to be contiguous.

### Versions

PyTorch version: 1.12.1+cu113
Is debug build: False
CUDA used to build PyTorch: 11.3
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.4 LTS (x86_64)
GCC version: (Ubuntu 9.4.0-1ubuntu1~20.04.1) 9.4.0
Clang version: Could not collect
CMake version: version 3.16.3
Libc version: glibc-2.31

Python version: 3.8.3 (default, Jul  2 2020, 16:21:59)  [GCC 7.3.0] (64-bit runtime)
Python platform: Linux-5.4.0-107-generic-x86_64-with-glibc2.10
Is CUDA available: True
CUDA runtime version: 11.2.152
CUDA_MODULE_LOADING set to: 
GPU models and configuration: 
GPU 0: NVIDIA RTX A6000
GPU 1: NVIDIA RTX A6000
GPU 2: NVIDIA RTX A6000
GPU 3: NVIDIA RTX A6000
GPU 4: NVIDIA RTX A6000
GPU 5: NVIDIA RTX A6000
GPU 6: NVIDIA RTX A6000
GPU 7: NVIDIA RTX A6000

Nvidia driver version: 470.103.01
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.8.4.0
/usr/lib/x86_64-linux-gnu/libcudnn_adv_infer.so.8.4.0
/usr/lib/x86_64-linux-gnu/libcudnn_adv_train.so.8.4.0
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_infer.so.8.4.0
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_train.so.8.4.0
/usr/lib/x86_64-linux-gnu/libcudnn_ops_infer.so.8.4.0
/usr/lib/x86_64-linux-gnu/libcudnn_ops_train.so.8.4.0
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] numpy==1.20.3
[pip3] torch==1.12.1+cu113
[pip3] torchaudio==0.12.1+cu113
[pip3] torchinfo==1.7.1
[pip3] torchvision==0.13.1+cu113
[pip3] torchviz==0.0.2
[conda] Could not collect


cc @VitalyFedyunin @jamesr66a