Unknown type name '__torch__.torch.classes.metal.Conv2dOpContext'

:bug: 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:

  1. I followed the tutorial to compile PyTorch from scratch.
  2. I download the HelloWorld example 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.
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
(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

Are you running the model on your phone, your mac or linux machine? cc @xta0

I’ve seen other people complain about this issue. Assigning this to myself.

I run it on a phone.