Skip to content

Commit a471305

Browse files
zou3519NicolasHug
andauthored
Update usages of torch.library APIs (pytorch#8384)
Co-authored-by: Nicolas Hug <[email protected]> Co-authored-by: Nicolas Hug <[email protected]>
1 parent e1f1fbe commit a471305

File tree

2 files changed

+2
-1
lines changed

2 files changed

+2
-1
lines changed

torchvision/_meta_registrations.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -160,7 +160,7 @@ def meta_ps_roi_pool_backward(
160160
return grad.new_empty((batch_size, channels, height, width))
161161

162162

163-
@torch._custom_ops.impl_abstract("torchvision::nms")
163+
@torch.library.register_fake("torchvision::nms")
164164
def meta_nms(dets, scores, iou_threshold):
165165
torch._check(dets.dim() == 2, lambda: f"boxes should be a 2d tensor, got {dets.dim()}D")
166166
torch._check(dets.size(1) == 4, lambda: f"boxes should have 4 elements in dimension 1, got {dets.size(1)}")

torchvision/csrc/ops/nms.cpp

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@ at::Tensor nms(
1919
}
2020

2121
TORCH_LIBRARY_FRAGMENT(torchvision, m) {
22+
m.set_python_module("torchvision._meta_registrations");
2223
m.def(TORCH_SELECTIVE_SCHEMA(
2324
"torchvision::nms(Tensor dets, Tensor scores, float iou_threshold) -> Tensor"));
2425
}

0 commit comments

Comments
 (0)