mmdeploy
mmdeploy copied to clipboard
Transfer to cpu only when torch.Tensor arguments in onnx export
Motivation
.cpu() in ONNX export inputs would be only supported in Tensor inputs so it should be only applied to Tensor types
Modification
Adds type checking before .cpu() and pass it directly if not.
Checklist
- Pre-commit or other linting tools are used to fix the potential lint issues.
- The modification is covered by complete unit tests. If not, please add more unit tests to ensure the correctness.
- If the modification has a dependency on downstream projects of a newer version, this PR should be tested with all supported versions of downstream projects.
- The documentation has been modified accordingly, like docstring or example tutorials.
If this MR doesn't get merged here, feel free to open it against our VBTI fork at https://github.com/VBTI-development/onedl-mmdeploy - we're actively maintaining updated versions with latest PyTorch support and would welcome the contribution!