Repository: pytorch
Issue: Unknown type name '__torch__.torch.classes.metal.Conv2dOpContext'
Labels: module: convolution,module: macos,oncall: mobile,module: ios

## 🐛 Bug

<!-- A clear and concise description of what the bug is. -->

Using PyTorch mobile with metal backend, Conv2dOpContext is unknown type name when loading the model in Xcode.

## To Reproduce

Steps to reproduce the behavior:

1. I followed the [tutorial](https://pytorch.org/tutorials/prototype/ios_gpu_workflow.html) to compile PyTorch from scratch.
2. I download the HelloWorld [example](https://github.com/pytorch/ios-demo-app/tree/master/HelloWorld) and change the `model.pt` to my own model.
3. My model with CPU backend runs normally but encounters the below error message with Metal backend.

<!-- If you have a code sample, error messages, stack traces, please provide it here as well -->

```
Unknown type name '__torch__.torch.classes.metal.Conv2dOpContext':
Serialized   File "code/__torch__/modules/model_name/___torch_mangle_270.py", line 5
  __buffers__ = []
  __annotations__ = []
  __annotations__["prepack_folding_forward._jit_pass_packed_weight_0"] = __torch__.torch.classes.metal.Conv2dOpContext
                                                                         ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
  __annotations__["prepack_folding_forward._jit_pass_packed_weight_1"] = __torch__.torch.classes.metal.Conv2dOpContext
  __annotations__["prepack_folding_forward._jit_pass_packed_weight_2"] = __torch__.torch.classes.metal.Conv2dOpContext
```

## Expected behavior

<!-- A clear and concise description of what you expected to happen. -->

## Environment

Please copy and paste the output from our
[environment collection script](https://raw.githubusercontent.com/pytorch/pytorch/master/torch/utils/collect_env.py)
(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.9.0a0+gitd69c22d
Is debug build: False
CUDA used to build PyTorch: Could not collect
ROCM used to build PyTorch: N/A
OS: KDE neon User Edition 5.22 (x86_64)
GCC version: (Ubuntu 9.3.0-17ubuntu1~20.04) 9.3.0
Clang version: Could not collect
CMake version: version 3.20.5
Libc version: glibc-2.31
Python version: 3.9.5 | packaged by conda-forge | (default, Jun 19 2021, 00:32:32)  [GCC 9.3.0] (64-bit runtime)
Python platform: Linux-5.11.0-22-generic-x86_64-with-glibc2.31
Is CUDA available: False
CUDA runtime version: Could not collect
GPU models and configuration: 
GPU 0: NVIDIA GeForce RTX 2080 Ti
GPU 1: NVIDIA GeForce RTX 2080 Ti
Nvidia driver version: 470.42.01
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.21.0
[pip3] torch==1.10.0a0+git512448a
[conda] magma-cuda113             2.5.2                         1    pytorch
[conda] mkl                       2021.2.0           h06a4308_296  
[conda] mkl-include               2021.3.0           h726a3e6_557    conda-forge
[conda] numpy                     1.21.0           py39hdbf815f_0    conda-forge
[conda] torch                     1.10.0a0+git512448a          pypi_0    pypi

## Additional context

<!-- Add any other context about the problem here. -->

Operators listed for `torch.jit.export_opnames(optimized_model)`
```
['aten::__and__.bool', 'aten::__getitem__.t', 'aten::_set_item.str', 'aten::add.int', 'aten::append.t', 'aten::arange', 'aten::avg_pool2d', 'aten::cat', 'aten::constant_pad_nd', 'aten::dim', 'aten::div.Scalar', 'aten::eq.int', 'aten::floor.float', 'aten::format', 'aten::gt.int', 'aten::le.int', 'aten::len.t', 'aten::lt.int', 'aten::mul.Scalar', 'aten::mul.Tensor', 'aten::ne.int', 'aten::ne.int_float', 'aten::pop.t', 'aten::repeat', 'aten::reshape', 'aten::size', 'aten::slice.t', 'aten::softmax.int', 'aten::sub.Scalar', 'aten::sub.int', 'aten::sum.dim_IntList', 'aten::to.prim_Device', 'aten::unsqueeze', 'aten::upsample_nearest1d.vec', 'aten::upsample_nearest2d.vec', 'aten::upsample_nearest3d.vec', 'aten::view', 'metal_prepack::conv2d_run', 'prim::RaiseException', 'prim::Uninitialized']
```

Got this message at the end of compilation.
```
warning: /Applications/Xcode.app/Contents/Developer/Toolchains/XcodeDefault.xctoolchain/usr/bin/ranlib: archive library: /path/build_ios/install/lib/libtorch.a the table of contents is empty (no object file members in the library define global symbols)
```

Running `grep` shows that `Conv2dOpContext` is not present in `libtorch.a`
```
(pytorch_metal) ➜  pytorch_metal git:(512448a425) ✗ grep "Conv2dOpContext" -r build_ios/install
build_ios/install/include/ATen/native/metal/ops/MetalConvolution.h:Tensor conv2d(const Tensor& input, Conv2dOpContext& context);
build_ios/install/include/ATen/native/metal/MetalPrepackOpContext.h:class Conv2dOpContext : public torch::jit::CustomClassHolder {
build_ios/install/include/ATen/native/metal/MetalPrepackOpContext.h:  Conv2dOpContext() = delete;
build_ios/install/include/ATen/native/metal/MetalPrepackOpContext.h:  Conv2dOpContext(
Binary file build_ios/install/lib/libtorch_cpu.a matches
```

cc @malfet @albanD @gmagogsfm