Errors when trying to run jit script

jit
1

I was trying to run this code section and with my limited understanding of jit, thought I can get a performance improvement on the function.

@torch.jit.script
def nms_pytorch2(dets, thresh):
    overlap = bbox_overlap(dets, dets)

    treshold_matrix = torch.tril((overlap > thresh), diagonal=-1)

    # Tensor elements indicate whether box should be kept
    is_maximum = treshold_matrix.new_ones(dets.shape[0])

    # loop over all boxes with highest confidence in the scene
    # Apply this vectorized over all boxes in the batch.
    for box in treshold_matrix.unbind(-1):
        # Disable all other boxes in the same scene if the current box is not
        # disabled.
        is_maximum = is_maximum & ~box

        # Also disable the overlaps of boxes which getting disabled right now.
        treshold_matrix &= ~box.unsqueeze(-2)

    return is_maximum

However, I got the following error and am not sure how to handle it

torch.jit.frontend.NotSupportedError: unsupported kind of augumented assignment: BitAnd:
         # Also disable the overlaps of boxes which getting disabled right now.
        treshold_matrix &= ~box.unsqueeze(-2)
                        ~~ <--- HERE
    return is_maximum

Your help would be really appreciated

2

Just figured the issue out, the operations need to be non in-place. Just putting a temp variable and assigning value of this temp variable in the next line solved the issue for me.