Documenting jit-compiled PyTorch Class Method (Sphinx)

Question

I am having a problem trying to document custom PyTorch models with Sphinx: methods that are jit-compiled show up without docstrings in the documentation. How do I fix this? I checkoed out Python Sphinx autodoc and decorated members and How to autodoc decorated methods with sphinx? but the proposed solutions don't seem to work. When I try using ..automethod I get

AttributeError: '_CachedForward' object has no attribute '__getattr__'

Here's a MWE; at the moment I circumvent the problem by writing a my_forward and calling it in the forward.

from torch import jit, Tensor


class DummyModel(jit.ScriptModule):
    """My dummy model"""

    def __init__(self, const: float):
        super(DummyModel, self).__init__()
        self.const = Tensor(const)

    @jit.script_method
    def forward(self, x: Tensor) -> Tensor:
        """This method does show as a :undoc-member: in the documentation"""
        return self.my_forward(x)

    def my_forward(self,  x: Tensor) -> Tensor:
        """This method shows up as a :member: in the documentation"""
        return x + self.const

Answer 1

Set the PYTORCH_JIT environment variable to 0 when running Sphinx. This will disable script and tracing annotations (decorators).

See https://pytorch.org/docs/stable/jit.html#disable-jit-for-debugging.