Torch.jit.script(module) vs @torch.jit.script装饰器写法区别

下面通过装饰器使用jit的写法是会报错的：

import torch

@torch.jit.script
class MyCell(torch.nn.Module):
    def __init__(self):
        super(MyCell, self).__init__()
        self.linear = torch.nn.Linear(4, 4)

    def forward(self, x, h):
        new_h = torch.tanh(self.linear(x) + h)
        return new_h, new_h

my_cell = MyCell()
x, h = torch.rand(3, 4), torch.rand(3, 4)
traced_cell = torch.jit.script(my_cell, (x, h))
print(traced_cell)
traced_cell(x, h)

"C:\Users\Administrator\AppData\Local\Packages\PythonSoftwareFoundation.Python.3.8_qbz5n2kfra8p0\LocalCache\local-packages\Python38\site-packages\torch\jit\__init__.py", line 1262, in script
    raise RuntimeError("Type '{}' cannot be compiled since it inherits"
RuntimeError: Type '<class '__main__.MyCell'>' cannot be compiled since it inherits from nn.Module, pass an instance instead

而下面的代码是正确的：

class MyCell(torch.nn.Module):
    def __init__(self):
        super(MyCell, self).__init__()
        self.linear = torch.nn.Linear(4, 4)

    def forward(self, x, h):
        new_h = torch.tanh(self.linear(x) + h)
        return new_h, new_h

my_cell = MyCell()
x, h = torch.rand(3, 4), torch.rand(3, 4)
traced_cell = torch.jit.script(my_cell, (x, h))
print(traced_cell)
traced_cell(x, h)

因为@torch.jit.script是用来装饰函数的，故用在类对象上面会报错，而下面两段代码是等价的：

@torch.jit.script
def fn(a: int):
   return a + 1

fn(3) # fn here is a scripted function

def fn(a: int):
   return a + 1

s_fn = torch.jit.script(fn)
s_fn(3) # s_fn here is a scripted function

参考Torch.jit.script(module) vs @torch.jit.script decorator - jit - PyTorch Forums