Allow to return `ONNXProgram` when calling `to_onnx(dynamo=True)`

[GdoongMathew]

Description & Motivation

Since torch 2.7.0, they started to unify the onnx export logic, thus when calling

torch.onnx.export(dynamo=True)

torch would start to return the ONNXProgram object. (pytorch/pytorch#137296)

The LightningModule should also have the same behavior.

Solution

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

Pitch

No response

Alternatives

No response

Additional context

No response

cc @lantiga @Borda

[GdoongMathew]

Related torch issue

pytorch/pytorch#151693