Repository: pytorch
Issue: Channels last performance problem
Labels: module: performance,triaged,module: memory format,Stale

### 🐛 Describe the bug

I tried to use channels last memory format as described in this guide: https://pytorch.org/tutorials/intermediate/memory_format_tutorial.html#converting-existing-models. 

However, it made my model 2x slower in Torchscript inference. My script is as follows:
import torch.nn as nn
import timeit
import torch
import matplotlib.pyplot as plt
import time

def time_gpu_model(model, shape: tuple) -> float:
    source = torch.rand((128, 3, *shape)).cuda().half().to(memory_format=torch.channels_last)
    reference = torch.rand((128, 3, *shape)).cuda().half().to(memory_format=torch.channels_last)
    with torch.no_grad():
        start = time.time()
        for i in range(10):
            model(source, reference)
        return time.time() - start


shapes = [(256, 256)]
model = torch.jit.load('trace_gpu.pt').half().to(memory_format=torch.channels_last)
time_gpu = list(map(lambda x: time_gpu_model(model, x), shapes))
print(time_gpu)
~             

Profiling, I find 60% of the time is spent on the following kernel:                              

_ZN2at6native27unrolled_elementwise_kernelIZZZNS0_21copy_device_to_deviceERNS_14TensorIteratorEbENKUlvE0_clEvENKUlvE18_clEvEUlN3c104HalfEE_NS_6detail5ArrayIPcLi2EEE16OffsetCalculatorILi1EjESE_NS0_6memory15LoadWithoutCastENSF_16StoreWithoutCastEEEviT_T0_T1_T2_T3_T4_

It was invoked 867 times. 

I understand this kernel is used to swap data formats. But shouldn't it just be done at the start?

### Versions


PyTorch version: 1.8.1+cu111
Is debug build: False
CUDA used to build PyTorch: 11.1
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.10

Python version: 3.7.10 | packaged by conda-forge | (default, Feb 19 2021, 16:07:37)  [GCC 9.3.0] (64-bit runtime)
Python platform: Linux-5.4.0-1060-aws-x86_64-with-debian-buster-sid
Is CUDA available: True
CUDA runtime version: 11.1.105
GPU models and configuration: GPU 0: Tesla T4
Nvidia driver version: 450.119.03
cuDNN version: Probably one of the following:
/usr/local/cuda-10.1/targets/x86_64-linux/lib/libcudnn.so.7.6.5
/usr/local/cuda-10.2/targets/x86_64-linux/lib/libcudnn.so.7.6.5
/usr/local/cuda-11.0/targets/x86_64-linux/lib/libcudnn.so.8.0.5
/usr/local/cuda-11.0/targets/x86_64-linux/lib/libcudnn_adv_infer.so.8.0.5
/usr/local/cuda-11.0/targets/x86_64-linux/lib/libcudnn_adv_train.so.8.0.5
/usr/local/cuda-11.0/targets/x86_64-linux/lib/libcudnn_cnn_infer.so.8.0.5
/usr/local/cuda-11.0/targets/x86_64-linux/lib/libcudnn_cnn_train.so.8.0.5
/usr/local/cuda-11.0/targets/x86_64-linux/lib/libcudnn_ops_infer.so.8.0.5
/usr/local/cuda-11.0/targets/x86_64-linux/lib/libcudnn_ops_train.so.8.0.5
/usr/local/cuda-11.1/targets/x86_64-linux/lib/libcudnn.so.8.0.5
/usr/local/cuda-11.1/targets/x86_64-linux/lib/libcudnn_adv_infer.so.8.0.5
/usr/local/cuda-11.1/targets/x86_64-linux/lib/libcudnn_adv_train.so.8.0.5
/usr/local/cuda-11.1/targets/x86_64-linux/lib/libcudnn_cnn_infer.so.8.0.5
/usr/local/cuda-11.1/targets/x86_64-linux/lib/libcudnn_cnn_train.so.8.0.5
/usr/local/cuda-11.1/targets/x86_64-linux/lib/libcudnn_ops_infer.so.8.0.5
/usr/local/cuda-11.1/targets/x86_64-linux/lib/libcudnn_ops_train.so.8.0.5
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] mypy-extensions==0.4.3
[pip3] numpy==1.19.2
[pip3] numpydoc==1.1.0
[pip3] torch==1.8.1+cu111
[pip3] torch-model-archiver==0.2.1
[pip3] torchserve==0.2.1
[pip3] torchvision==0.9.1+cu111
[conda] blas                      1.0                         mkl  
[conda] mkl                       2020.2                      256  
[conda] mkl-service               2.3.0            py37he8ac12f_0  
[conda] mkl_fft                   1.2.1            py37h54f3939_0  
[conda] mkl_random                1.1.1            py37h0573a6f_0  
[conda] mypy_extensions           0.4.3                    py37_0  
[conda] numpy                     1.19.2           py37h54aff64_0  
[conda] numpy-base                1.19.2           py37hfa32c7d_0  
[conda] numpydoc                  1.1.0              pyhd3eb1b0_1  
[conda] torch                     1.8.1+cu111              pypi_0    pypi
[conda] torch-model-archiver      0.2.1                    py37_0    pytorch
[conda] torchserve                0.2.1                    py37_0    pytorch
[conda] torchvision               0.9.1+cu111              pypi_0    pypi


cc @VitalyFedyunin @ngimel @jamesr66a