Model not compatible with TorchScript conversion (via torch.jit.script) about movinet-pytorch HOT 4 OPEN

Unfortunately I have no plans to support TorchScript at the moment. I'm also not convinced that's the only problem, for example the computation of same padding could raise an error. If that's the case I suspect it would also be necessary to train the models from scratch in pytorch.

from movinet-pytorch.

The type hint issue was easy to get around but even more problematic is the use of einops which appears to be noncompat with torchscript

from torch import nn
import torch
from einops import rearrange
class Foo(nn.Module):
    def __init__(self):
        super(Foo, self).__init__()
    def forward(self, x):
        return rearrange(x, 'a b -> b a')
torch.jit.script(Foo())
...
NotSupportedError: Compiled functions can't take variable number of arguments or use keyword-only arguments with defaults:
  File "/Users/sean/standard/.env/lib/python3.6/site-packages/einops/einops.py", line 393
def rearrange(tensor, pattern: str, **axes_lengths):
                                     ~~~~~~~~~~~~~ <--- HERE

It looks like tracing may be compatible, though is more restrictive, per the issue on einops arogozhnikov/einops#115

foo = Foo()
traced_foo = torch.jit.trace(foo, torch.rand(3, 3))

from movinet-pytorch.

Is there any solution to this problem?

from movinet-pytorch.

I have a fork that modifies the models so that it can be exported to TorchScript or ONNX.
https://github.com/Subalzero/MoViNet-pytorch

from movinet-pytorch.