Repository: pytorch
Issue: Error while converting model for NNAPI "Unsupported node kind ('quantized::batch_norm2d')"
Labels: triaged,module: android,oncall: mobile

## 🐛 Bug

I'm converting the Unet model from (torch-hub)[https://pytorch.org/hub/mateuszbuda_brain-segmentation-pytorch_unet/] to run on Android for a benchmark. Error happens after I've quantized the model and converted to use `NHWC`. 

## To Reproduce

Steps to reproduce the behavior:

```python
model = torch.hub.load('mateuszbuda/brain-segmentation-pytorch', 'unet',
    in_channels=3, out_channels=1, init_features=32, pretrained=True)
model.eval()

# set quant back-end
backend = 'qnnpack'

# Set the Quant backend
model.qconfig = torch.quantization.get_default_qconfig(backend)
torch.backends.quantized.engine = backend

# prepare watchers and initialize the quant stubs
static_quant_model = torch.quantization.prepare(model)

# Do a forward with dummy data to init the scales
dummy = prepare(image)
static_quant_model(dummy) #forward pass

# convert to int8 quantized model
model_int8 = torch.quantization.convert(static_quant_model)

#logic to cast input to correct qint type
input_tensor = quantizer(dummy)
input_tensor = input_tensor.contiguous(memory_format=torch.channels_last)
input_tensor.nnapi_nhwc = True


with torch.no_grad():
    nhwc_int8 = torch.jit.trace(model_int8, input_tensor)

nnapi_model = torch.backends._nnapi.prepare.convert_model_to_nnapi(nhwc_int8, input_tensor)
```

Results in:
```
---------------------------------------------------------------------------
Exception                                 Traceback (most recent call last)
<ipython-input-48-88c17d678fa3> in <module>
----> 1 nnapi_model = torch.backends._nnapi.prepare.convert_model_to_nnapi(nhwc_int8, input_tensor)

envs/py3/lib/python3.6/site-packages/torch/backends/_nnapi/prepare.py in convert_model_to_nnapi(model, inputs)
    167         delist_outputs = False
    168 
--> 169     ser_model, used_weights, inp_mem_fmts, out_mem_fmts = serialize_model(model, inputs)
    170     ser_model_tensor = torch.tensor(list(ser_model), dtype=torch.uint8)
    171 

envs/py3/lib/python3.6/site-packages/torch/backends/_nnapi/serializer.py in serialize_model(module, inputs, config)
   1361 
   1362 def serialize_model(module, inputs, config=None):
-> 1363     return _NnapiSerializer(config).serialize_model(module, inputs)

envs/py3/lib/python3.6/site-packages/torch/backends/_nnapi/serializer.py in serialize_model(self, model, inputs)
    518         for idx, node in enumerate(model.graph.nodes()):
    519             LOG.debug("Processing node #%d: %r", idx, node)
--> 520             self.add_node(node)
    521 
    522         retn = model.graph.return_node()

envs/py3/lib/python3.6/site-packages/torch/backends/_nnapi/serializer.py in add_node(self, node)
    639         adder = self.ADDER_MAP.get(node.kind())
    640         if not adder:
--> 641             raise Exception("Unsupported node kind (%r) in node %r" % (node.kind(), node))
    642         adder(self, node)
    643 

Exception: Unsupported node kind ('quantized::batch_norm2d') in node %input.3 : Tensor = quantized::batch_norm2d(%input.2, %self.encoder1.enc1norm1.weight, %self.encoder1.enc1norm1.bias, %self.encoder1.enc1norm1.running_mean, %self.encoder1.enc1norm1.running_var, %55, %56, %57), scope: __module.encoder1/__module.encoder1.enc1norm1 # envs/py3/lib/python3.6/site-packages/torch/nn/quantized/modules/batchnorm.py:17:0
```

## Expected behavior

Model is converted to NNAPI compliant version.

## Environment

```
Collecting environment information...
PyTorch version: 1.8.1+cu102
Is debug build: False
CUDA used to build PyTorch: 10.2
ROCM used to build PyTorch: N/A

OS: Ubuntu 18.04.3 LTS (x86_64)
GCC version: (Ubuntu 7.5.0-3ubuntu1~18.04) 7.5.0
Clang version: Could not collect
CMake version: version 3.14.0

Python version: 3.6 (64-bit runtime)
Is CUDA available: True
CUDA runtime version: Could not collect
GPU models and configuration: 
GPU 0: GeForce RTX 2080 Ti
GPU 1: GeForce RTX 2080 Ti
GPU 2: GeForce RTX 2080 Ti
GPU 3: GeForce RTX 2080 Ti

Nvidia driver version: 440.100
cuDNN version: /usr/lib/x86_64-linux-gnu/libcudnn.so.7.6.5
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] numpy==1.19.5
[pip3] pytorch-lightning==1.2.4
[pip3] pytorch-metric-learning==0.9.98
[pip3] pytorch-msssim==0.2.1
[pip3] torch==1.8.1
[pip3] torch-optimizer==0.0.1a9
[pip3] torchaudio==0.8.1
[pip3] torchmetrics==0.2.0
[pip3] torchtext==0.9.1
[pip3] torchvision==0.9.1
[conda] _pytorch_select           0.2                       gpu_0  
[conda] blas                      1.0                         mkl    conda-forge
[conda] cudatoolkit               10.0.130                      0  
[conda] libblas                   3.8.0                    14_mkl    conda-forge
[conda] libcblas                  3.8.0                    14_mkl    conda-forge
[conda] liblapack                 3.8.0                    14_mkl    conda-forge
[conda] magma-cuda100             2.5.2                         1    pytorch
[conda] mkl                       2019.4                      243  
[conda] mkl-include               2020.2                      256  
[conda] mkl-service               2.3.0            py36he904b0f_0  
[conda] mkl_fft                   1.0.15           py36ha843d7b_0  
[conda] mkl_random                1.1.0            py36hd6b4f25_0  
[conda] numpy                     1.19.5                   pypi_0    pypi
[conda] pytorch-lightning         1.2.4                    pypi_0    pypi
[conda] pytorch-metric-learning   0.9.98                   pypi_0    pypi
[conda] pytorch-msssim            0.2.1                    pypi_0    pypi
[conda] torch                     1.8.1                    pypi_0    pypi
[conda] torch-optimizer           0.0.1a9                  pypi_0    pypi
[conda] torchaudio                0.8.1                    pypi_0    pypi
[conda] torchmetrics              0.2.0                    pypi_0    pypi
[conda] torchtext                 0.9.1                    pypi_0    pypi
[conda] torchvision               0.9.1                    pypi_0    pypi
```
## Additional context

<!-- Add any other context about the problem here. -->