Почему я не могу использовать классы для аннотации типов аргументов функции в декораторе `torch.jit.script`?

Question

Этот код компилируется отлично:

import torch
import torch.nn             as nn

class Foo(nn.Module):
    def __init__(self):
        super(Foo, self).__init__()
        self.x = 0

    def forward(self, X):
        X      *= self.x
        self.x += 1
        return X

# @torch.jit.script
def bar(f: Foo):
    return f.x

Но, если я раскомментирую строку # @torch.jit.script, я получаю следующую ошибку:

Traceback (most recent call last):
  File "test1.py", line 18, in <module>
    def bar(f: Foo):
  File "/anaconda3/envs/pytorch/lib/python3.7/site-packages/torch/jit/__init__.py", line 1103, in script
    fn = torch._C._jit_script_compile(qualified_name, ast, _rcb, get_default_args(obj))
RuntimeError:
Unknown type name 'Foo':
at test1.py:18:12
@torch.jit.script
def bar(f: Foo):
           ~~~ <--- HERE
    return f.x

Если я изменю аннотацию типа на int:

@torch.jit.script
# def bar(f: Foo):
#     return f.x
def bar(f: int):
    return f

затем компиляция снова работает.

Кто-нибудь знает, что мне нужно сделать, чтобы позволить использовать определения моих пользовательских классов в аннотациях типов к аргументам функций, которые l ie под torch.jit.script декоратор?

Answer 1

Только список типов в документации здесь может использоваться в качестве аргументов функций:

https://pytorch.org/docs/stable/jit_language_reference.html#supported -type

nn.Module s имеют особую обработку в TorchScript, чтобы они работали, но в настоящее время они не поддерживаются в качестве аргументов.