Incompatible LightningModule.to_onnx type-hint

[dominicgkerr]

Bug description

Hello!

The LightningModule.to_onnx [1] method provides a file_path: Union[str, Path] type-hint. However, file_path is later passed (unmodified) into torch.onnx.export [2] which expects an Union[str, io.BytesIO] file handle (though, I think [3] suggests it should in fact only be a str...)

Personally, I'd like to be able to pass Path instances into LightningModule.to_onnx, which can be achieved by changing [4] to:

torch.onnx.export(self, input_sample, str(file_path), **kwargs)

Alternatively, removing the Path type-hint corrects and possible confusion, and wouldn't (as far as I can see) adversely effect any existing Lightning code...

I'd be very happy to open a small PR (let me know if you'd prefer to cast file_path to a str, or simply drop the Path type-hint)! Many thanks

• [1] https://github.com/Lightning-AI/pytorch-lightning/blob/master/src/lightning/pytorch/core/module.py#L1357
• [2] https://github.com/pytorch/pytorch/blob/main/torch/onnx/utils.py#L193
• [3] https://github.com/pytorch/pytorch/blob/main/torch/onnx/utils.py#L1603
• [4] https://github.com/Lightning-AI/pytorch-lightning/blob/master/src/lightning/pytorch/core/module.py#L1398

What version are you seeing the problem on?

master

How to reproduce the bug

from pathlib import Path
from lightning.pytorch.demos.boring_classes import BoringModel

model = BoringModel()

file_path = Path("./model.onnx")
model.to_onnx(file_path)

[awaelchli]

Thanks! Calling str before passing it to onnx sounds good to me.