Unverified Commit a4713054 authored by Richard Zou's avatar Richard Zou Committed by GitHub
Browse files

Update usages of torch.library APIs (#8384)


Co-authored-by: default avatarNicolas Hug <nh.nicolas.hug@gmail.com>
Co-authored-by: default avatarNicolas Hug <contact@nicolas-hug.com>
parent e1f1fbe4
......@@ -160,7 +160,7 @@ def meta_ps_roi_pool_backward(
return grad.new_empty((batch_size, channels, height, width))
@torch._custom_ops.impl_abstract("torchvision::nms")
@torch.library.register_fake("torchvision::nms")
def meta_nms(dets, scores, iou_threshold):
torch._check(dets.dim() == 2, lambda: f"boxes should be a 2d tensor, got {dets.dim()}D")
torch._check(dets.size(1) == 4, lambda: f"boxes should have 4 elements in dimension 1, got {dets.size(1)}")
......
......@@ -19,6 +19,7 @@ at::Tensor nms(
}
TORCH_LIBRARY_FRAGMENT(torchvision, m) {
m.set_python_module("torchvision._meta_registrations");
m.def(TORCH_SELECTIVE_SCHEMA(
"torchvision::nms(Tensor dets, Tensor scores, float iou_threshold) -> Tensor"));
}
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment