NameError: name '_C' is not defined: (Prototype) Use iOS GPU in PyTorch (Metal)

Hello, I wanted to ask a simple question just before I try PyTorch Metal integration.
In tutorial there is a sentence:
"Next we need to make some changes in TorchModule.mm"
Here should be a snippet of code, I guess, that shows how metal() should be used, because next sentence is:
“As you can see, we simply just call .metal() to move our input tensor from CPU to GPU”
So my question is - I should call .metal() on tensor argument just before calling forward method as I call .cuda() for example in Python?

@xta0 Is it something you know?

yes, call .metal() on your input tensor and call .cpu() on the output tensor to move it back. Something like

auto input = torch.rand({1,3,2,2}).metal()
auto output = model.forward({input}).toTensor().cpu();

Thank you, @xta0, but I struggle to import Metal PyTorch torch==1.9.0a0+gite442d5c
(after fixing the compilation error)
Error near import torch when I try to run my model:

Traceback (most recent call last):
  File "/demo.py", line 14, in <module>
    import torch
  File "/opt/anaconda3/envs/FOM_METAL/lib/python3.8/site-packages/torch/__init__.py", line 228, in <module>
    __all__ += [name for name in dir(_C)
NameError: name '_C' is not defined

Error in python console:

python
>>> import torch
  File "/opt/anaconda3/envs/FOM_METAL/lib/python3.8/site-packages/torch/_tensor.py", line 322, in Tensor
    detach = _C._add_docstr(_C._TensorBase.detach, r"""
RuntimeError: method 'detach' already has a docstring

torch 1.8.0 works fine

I have no clue how to fix this.

Well I’ve managed to compile it and even import, but now I’m getting some vague error when I try to optimize for Metal:

traced_fom = torch.jit.trace(fom, (source, driving, initial), strict=False)
optimized_fom = mob_op.optimize_for_mobile(traced_fom, backend='metal') # <<< error here
optimized_fom.save('fom-metal.pt')
Traceback (most recent call last):
  File "/demo.py", line 326, in <module>
    predictions = make_animation(source_image, driving_video)
  File "/demo.py", line 91, in make_animation
    optimized_fom = mob_op.optimize_for_mobile(traced_fom, backend='metal')
  File "/opt/anaconda3/envs/fom/lib/python3.8/site-packages/torch/utils/mobile_optimizer.py", line 69, in optimize_for_mobile
    optimized_cpp_module = torch._C._jit_pass_metal_optimize_for_mobile(script_module._c, preserved_methods_str)
RuntimeError: 0 INTERNAL ASSERT FAILED at "../torch/csrc/jit/ir/alias_analysis.cpp":532, 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, None, int[], int[], int[], int, None, None, 

We don’t have an op for metal_prepack::conv2d_prepack but it isn’t a special case. Argument types: Tensor, None, int[], int[], int[], int, None, None,

My parameters are:

conda 4.9.2
Python 3.6.13 :: Anaconda, Inc.

python -m pip freeze

attr==0.3.1
attrs==20.3.0
certifi==2020.12.5
cffi==1.11.5
cloudpickle==0.5.3
coremltools==4.1
cycler==0.10.0
Cython==0.29.22
dask==2021.3.0
dataclasses==0.8
decorator==4.3.0
imageio==2.3.0
imageio-ffmpeg==0.4.3
kiwisolver==1.0.1
matplotlib==2.2.2
mpmath==1.2.1
networkx==2.1
numpy==1.18.1
packaging==20.9
pandas==0.23.4
Pillow==5.2.0
protobuf==3.15.6
pycparser==2.18
pygit==0.1
pyparsing==2.2.0
python-dateutil==2.7.3
pytz==2018.5
PyWavelets==0.5.2
PyYAML==5.1
scikit-image==0.15.0
scikit-learn==0.19.2
scikit-video==1.1.11
scipy==1.1.0
six==1.11.0
sympy==1.7.1
toolz==0.9.0
torch==1.9.0a0+gite442d5c
torchvision==0.9.0
tqdm==4.24.0
typing-extensions==3.7.4.3

Should I use some specific torch version to make it work?
Related bug report