"Unknown type name": forward reference breaking torchscript

jit
1

I am trying to script a model with the zuko package (GitHub - probabilists/zuko: Normalizing flows in PyTorch). I instantiate a model as
flow = zuko.flows.MAF(features=1, context=1, transforms=4, hidden_features=[32,32])

then after training, I script it as
sm = torch.jit.script(flow)

which results in the error

Unknown type name 'LazyTransform':
  File "/global/homes/r/rmastand/.local/lib/python3.8/site-packages/zuko/lazy.py", line 74
    @property
    def inv(self) -> LazyTransform:
                     ~~~~~~~~~~~~~ <--- HERE
        r"""A lazy inverse transformation :math:`x = f^{-1}(y | c)`."""

The issue comes from [this file] (zuko/zuko/lazy.py at 313481e0873b4c98636314d6c3feaea30284b28a · probabilists/zuko · GitHub) — LazyTransform is not defined until line 80, but it is forward-referred to on line 74. Is there a decorator that I can attach to these functions? The entire repository seems based around the LazyTransform, so I don’t think I can avoid compiling it.