Closed qihqi closed 3 weeks ago
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
🚀 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