Bug

Using PyTorch mobile with metal backend, Conv2dOpContext is unknown type name when loading the model in Xcode.

To Reproduce
Steps to reproduce the behavior:

I followed the tutorial 2 to compile PyTorch from scratch.
I download the HelloWorld example and change the model.pt to my own model.
My model with CPU backend runs normally but encounters the below error message with Metal backend.


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


Environment
Please copy and paste the output from our
environment collection script 2
(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

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 the 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