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 8