torch.onnx.export(dynamo=True)

class LightningModule:
    @torch.no_grad()
    def to_onnx(
        self,
        file_path: Union[str, Path, BytesIO, None] = None,
        input_sample: Optional[Any] = None,
        **kwargs: Any,
    ) -> Union["ONNXProgram", None]:
        if kwargs.get("dynamo", False) and not _ONNXSCRIPT_AVAILABLE:
            raise ModuleNotFoundError(...)
        ...
        ret = torch.onnx.export(self, input_sample, file_path, **kwargs)
        self.train(mode)
        return ret