Error running jit model on medal, any possible solutions?

username: Matt_Davis
cooked: <p>Open to ideas on if this is an actual missing issue or something I can solve.</p>
<p>RuntimeError                              Traceback (most recent call last)<br>
 in ()<br>
1 kp_driving = kp_detector(image_tensor)<br>
2 traced_generator = torch.jit.trace(generator, (image_tensor, kp_driving, kp_driving), strict=False)<br>
----&gt; 3 optimized_traced_generator = optimize_for_mobile(traced_generator, backend=‘metal’)<br>
4 optimized_traced_generator._save_for_lite_interpreter("/content/generator_metal.pt")</p>
<p>/usr/local/lib/python3.7/dist-packages/torch/utils/mobile_optimizer.py in optimize_for_mobile(script_module, optimization_blocklist, preserved_methods, backend)<br>
67         optimized_cpp_module = torch._C._jit_pass_vulkan_optimize_for_mobile(script_module._c, preserved_methods_str)<br>
68     elif backend == ‘metal’:<br>
—&gt; 69         optimized_cpp_module = torch._C._jit_pass_metal_optimize_for_mobile(script_module._c, preserved_methods_str)<br>
70     else:<br>
71         raise TypeError(“Unknown backend, must be one of ‘CPU’, ‘Vulkan’ or ‘Metal’”)</p>
<p>RuntimeError: 0INTERNAL ASSERT FAILED at “…/torch/csrc/jit/ir/alias_analysis.cpp”:584, please report a bug to PyTorch. We don’t have an op for metal_prepack::conv2d_prepack but it isn’t a special case.  Argument types: Tensor, Tensor, int[], int[], int[], int, NoneType, NoneType,</p>

username: ptrblck
cooked: <p>Based on the error message it sounds like an unexpected bug, so could you create a <a href="https://github.com/pytorch/pytorch/issues">GitHub issue</a> so that we can track and fix it, please?</p>

username: Matt_Davis
cooked: <p>Is it possible to export metal model on linux? experiencing a similar issue.</p>

username: ptrblck
cooked: <p>Yes, I think exporting the model should work on different platform, but the execution would then need the Metal backend. Are you seeing another issue with it, since you’ve created the topic?<br>
Did you create an issue on GitHub, so that the devs could take a look at the issue?</p>

username: Matt_Davis
cooked: <p>A collegue opened one up.</p><aside class="onebox githubissue" data-onebox-src="https://github.com/pytorch/pytorch/issues/69609">
  <header class="source">

      <a href="https://github.com/pytorch/pytorch/issues/69609" target="_blank" rel="noopener nofollow ugc">github.com/pytorch/pytorch</a>
  </header>

  <article class="onebox-body">
    <div class="github-row">
  <div class="github-icon-container" title="Issue">
	  <svg width="60" height="60" class="github-icon" viewbox="0 0 14 16" aria-hidden="true"><path d="M7 2.3c3.14 0 5.7 2.56 5.7 5.7s-2.56 5.7-5.7 5.7A5.71 5.71 0 0 1 1.3 8c0-3.14 2.56-5.7 5.7-5.7zM7 1C3.14 1 0 4.14 0 8s3.14 7 7 7 7-3.14 7-7-3.14-7-7-7zm1 3H6v5h2V4zm0 6H6v2h2v-2z"></path></svg>
  </div>

  <div class="github-info-container">
    <h4>
      <a href="https://github.com/pytorch/pytorch/issues/69609" target="_blank" rel="noopener nofollow ugc">[mobile_optimizer] RuntimeError: We don't have an op for metal_prepack::conv2d_prepack</a>
    </h4>

    <div class="github-info">
      <div class="date">
        opened <span class="discourse-local-date" data-format="ll" data-date="2021-12-08" data-time="13:36:18" data-timezone="UTC">01:36PM - 08 Dec 21 UTC</span>
      </div>


      <div class="user">
        <a href="https://github.com/anhappdev" target="_blank" rel="noopener nofollow ugc">
          <img alt="anhappdev" src="https://avatars.githubusercontent.com/u/85728587?v=4" class="onebox-avatar-inline" width="20" height="20">
          anhappdev
        </a>
      </div>
    </div>

    <div class="labels">
        <span style="display:inline-block;margin-top:2px;background-color: #B8B8B8;padding: 2px;border-radius: 4px;color: #fff;margin-left: 3px;">
          oncall: mobile
        </span>
    </div>
  </div>
</div>

  <div class="github-row">
    <p class="github-body-container">## 🐛 Bug

When trying to optimize a PyTorch model for running on iOS devices w<span class="show-more-container"><a href="" rel="noopener" class="show-more">…</a></span><span class="excerpt hidden">ith Metal backend
```
traced_detector = torch.jit.trace(kp_detector, image_tensor, strict=False)
optimized_traced_detector = optimize_for_mobile(traced_detector, backend='metal')
```
I got the error:
```
RuntimeError: 0INTERNAL ASSERT FAILED at "../torch/csrc/jit/ir/alias_analysis.cpp":584, please report a bug to PyTorch. We don't have an op for metal_prepack::conv2d_prepack but it isn't a special case.  Argument types: Tensor, NoneType, int[], int[], int[], int, NoneType, NoneType, 
```

## To Reproduce

See [this Colab notebook](https://colab.research.google.com/drive/1X40mGa652BiJjVW_rT2uPBvcD2UAxeSa?usp=sharing) for steps to reproduce the error.
Here is the error for quick reference and text search:
```
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
&lt;ipython-input-18-6a416810c207&gt; in &lt;module&gt;()
      1 kp_driving = kp_detector(image_tensor)
      2 traced_generator = torch.jit.trace(generator, (image_tensor, kp_driving, kp_driving), strict=False)
----&gt; 3 optimized_traced_generator = optimize_for_mobile(traced_generator, backend='metal')
      4 optimized_traced_generator._save_for_lite_interpreter("/content/generator_metal.pt")

/usr/local/lib/python3.7/dist-packages/torch/utils/mobile_optimizer.py in optimize_for_mobile(script_module, optimization_blocklist, preserved_methods, backend)
     67         optimized_cpp_module = torch._C._jit_pass_vulkan_optimize_for_mobile(script_module._c, preserved_methods_str)
     68     elif backend == 'metal':
---&gt; 69         optimized_cpp_module = torch._C._jit_pass_metal_optimize_for_mobile(script_module._c, preserved_methods_str)
     70     else:
     71         raise TypeError("Unknown backend, must be one of 'CPU', 'Vulkan' or 'Metal'")

RuntimeError: 0INTERNAL ASSERT FAILED at "../torch/csrc/jit/ir/alias_analysis.cpp":584, please report a bug to PyTorch. We don't have an op for metal_prepack::conv2d_prepack but it isn't a special case.  Argument types: Tensor, Tensor, int[], int[], int[], int, NoneType, NoneType, 
```

## Expected behavior

Code run without error.

## Environment

```
Collecting environment information...
PyTorch version: 1.10.0+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: 6.0.0-1ubuntu2 (tags/RELEASE_600/final)
CMake version: version 3.12.0
Libc version: glibc-2.26

Python version: 3.7.12 (default, Sep 10 2021, 00:21:48)  [GCC 7.5.0] (64-bit runtime)
Python platform: Linux-5.4.104+-x86_64-with-Ubuntu-18.04-bionic
Is CUDA available: False
CUDA runtime version: 11.1.105
GPU models and configuration: Could not collect
Nvidia driver version: Could not collect
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.7.6.5
/usr/lib/x86_64-linux-gnu/libcudnn.so.8.0.5
/usr/lib/x86_64-linux-gnu/libcudnn_adv_infer.so.8.0.5
/usr/lib/x86_64-linux-gnu/libcudnn_adv_train.so.8.0.5
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_infer.so.8.0.5
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_train.so.8.0.5
/usr/lib/x86_64-linux-gnu/libcudnn_ops_infer.so.8.0.5
/usr/lib/x86_64-linux-gnu/libcudnn_ops_train.so.8.0.5
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] msgpack-numpy==0.4.7.1
[pip3] numpy==1.19.5
[pip3] torch==1.10.0+cu111
[pip3] torchaudio==0.10.0+cu111
[pip3] torchsummary==1.5.1
[pip3] torchtext==0.11.0
[pip3] torchvision==0.11.1+cu111
[conda] Could not collect
```

## Additional context

There are 2 closed issues with same error. But trying what they stated did not work (use nightly build or save/load model from file)
https://github.com/pytorch/pytorch/issues/54377
https://github.com/pytorch/pytorch/issues/66307</span></p>
  </div>

  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>

<p>Let me know your thoughts.</p>
<p>Matt</p>

username: Matt_Davis
cooked: <p>Just looking at this more closely.</p>
<p>Is there a case sensitive issue or just no unit test around this for the changes to the metal to conv2d prepack? Seems like it is supported, just not working, maybe a previous version works.</p>
<p>I see the support implementation in PRs like <a href="https://github.com/pytorch/pytorch/commit/8ae8fb7dd11ebc3f35dcf17b639c0bb47b9ff235" class="inline-onebox" rel="noopener nofollow ugc">[iOS GPU][Stub] Move conv2d_prepack impl from MetalPrepackOpRegister.… · pytorch/pytorch@8ae8fb7 · GitHub</a></p>
<p>The error seems to be the error is<br>
/usr/local/lib/python3.7/dist-packages/torch/utils/mobile_optimizer.py in optimize_for_mobile(script_module, optimization_blocklist, preserved_methods, backend)<br>
67         optimized_cpp_module = torch._C._jit_pass_vulkan_optimize_for_mobile(script_module._c, preserved_methods_str)<br>
68     elif backend == ‘metal’:<br>
—&gt; 69         optimized_cpp_module = torch._C._jit_pass_metal_optimize_for_mobile(script_module._c, preserved_methods_str)<br>
70     else:<br>
71         raise TypeError(“Unknown backend, must be one of ‘CPU’, ‘Vulkan’ or ‘Metal’”)</p>
<p>Seems like a small compilation issue. I tried switching the label on export to ‘Metal’ no affect</p>