So you have the IoU part up to ovr, which is “item by item” this typically can be sped up by combining into a single kernel. As this type of operation is typically memory bandwidth bound, so the speedup factor is roughly equal to the number of kernel calls (function calls) you had before. You could adapt that into returning a boolean mask. This could, however, be also automatically fused in the JIT, so it might not be worth doing C++ coding for that if that’s not your favorite thing to do.
I have a notebook on this (using IoU as an example) in the JIT part of my PyTorch, JIT, Android talk 1 from December 2018 - a classic if you want.
The thresholding (which I’d write by boolean indexing returned from the JITed function rather than where/nonzero) is also largely done in the TorchVision kernel, doing this externally adds some overhead, but in the overall scheme of 3d convolutions that likely come before, it might not be that significant.
Best regards
Thomas