Fix tv_tensors.wrap to preserve subclass types for BoundingBoxes and KeyPoints#9332
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/vision/9332
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit 1dfd61b with merge base 7a78e54 ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
|
@raimbekovm Thank you for the PR. Could you please add the non-regression test to test your implementation? Feel free to let me know if you need more guidance for that. |
Test that tv_tensors.wrap() preserves subclass types for BoundingBoxes and KeyPoints instead of returning base class instances.
|
Done, added the regression test in 1ef69bf |
NicolasHug
left a comment
There was a problem hiding this comment.
Thanks for the PR @raimbekovm and @zy1git for the review!
|
Hey @NicolasHug! You merged this PR, but no labels were added. |
|
Thank you for the review! |
Summary
Fixes
tv_tensors.wrap()to preserve subclass types when wrapping instances of custom subclasses ofBoundingBoxesorKeyPoints.The bug: The function used hardcoded class names (
BoundingBoxes._wrap()andKeyPoints._wrap()) instead oftype(like)._wrap(), causing it to return base class instances instead of the actual subclass type.The fix: Replace
BoundingBoxes._wrap()withtype(like)._wrap()and similarly forKeyPoints._wrap(). This is consistent with how other TV tensor types are already handled (line 39).Fixes #9328
cc @vfdev-5