Skip to content

Commit

Permalink
Update torch compile
Browse files Browse the repository at this point in the history
  • Loading branch information
ericup committed Feb 8, 2024
1 parent 236fc00 commit 13255d4
Showing 1 changed file with 5 additions and 5 deletions.
10 changes: 5 additions & 5 deletions celldetection/ops/boxes.py
Expand Up @@ -10,16 +10,16 @@
__all__ = ['nms', 'contours2boxes', 'pairwise_box_iou', 'pairwise_generalized_box_iou', 'filter_by_box_voting']


def no_decorator(*args, **kwargs):
def torch_compile(*args, **kwargs):
def decorator(func):
return func
if WINDOWS:
return func # compile not supported on Windows, yet
else:
return torch.compile(func, *args, **kwargs)

return decorator


torch_compile = (no_decorator if WINDOWS else torch.compile) # TODO: Remove when torch.compile is supports on Windows


@torch_compile(dynamic=True)
def nms(boxes, scores, thresh=.5) -> torch.Tensor:
"""Non-maximum suppression.
Expand Down

0 comments on commit 13255d4

Please sign in to comment.