Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Implement torchvision.ops.nms in torch_xla2. #8278

Open
qihqi opened this issue Oct 18, 2024 · 0 comments · May be fixed by #8291
Open

Implement torchvision.ops.nms in torch_xla2. #8278

qihqi opened this issue Oct 18, 2024 · 0 comments · May be fixed by #8291
Assignees
@qihqi
Labels
torchxla2

Comments

@qihqi
Copy link
Collaborator

qihqi commented Oct 18, 2024

🚀 Feature

nms: https://pytorch.org/vision/stable/generated/torchvision.ops.nms.html
is an op used by many torchvision models.
It would be nice for those models to run on torchxla2.

Here is an implementation of it in Jax from MLperf submission:
https://raw.githubusercontent.com/mlperf/training_results_v0.7/refs/heads/master/Google/benchmarks/ssd/implementations/ssd-research-JAX-tpu-v3-4096/nms.py
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Assignees

@qihqi qihqi

Labels
torchxla2
Projects
None yet
Milestone
No milestone
Development

Successfully merging a pull request may close this issue.

Add impl for torchvision.nms and support multiple batch dim for convolution. pytorch/xla
1 participant
@qihqi