Skip to content

Commit

Permalink
Update usages of torch.library APIs
Browse files Browse the repository at this point in the history
We deprecated impl_abstract. This PR replaces it with the new API
(register_fake). register_fake also (sometimes) requires a
`set_python_module` in C++, so I add that as well.

Test Plan:
- existing tests
  • Loading branch information
zou3519 committed Apr 22, 2024
1 parent 96640af commit 47fa031
Show file tree
Hide file tree
Showing 2 changed files with 2 additions and 1 deletion.
2 changes: 1 addition & 1 deletion torchvision/_meta_registrations.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)}")
Expand Down
1 change: 1 addition & 0 deletions torchvision/csrc/ops/nms.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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"));
}
Expand Down

0 comments on commit 47fa031

Please sign in to comment.