@ptrblck Thank you so much for your feedback. I managed to reduce the function size further.
def nms_pytorch(dets, thresh):
    '''
    dets is a numpy array : num_dets, 6
    The detections are already in sorted order and so can be used directly.
    '''

    keep = []
    selected = torch.ones(dets.size(0))
    for i, bbox in enumerate(dets):
        if selected[i] == 0:
            continue
        keep.append(i)
        ovr = bbox_overlap(bbox.unsqueeze(0), dets[i+1:])
        # ovr is 1 X N tensor and we need to check for columns where it is greater than thresh
        inds = torch.where(ovr > thresh)[1]
        # We basically start from the first index and hence an offset has to be added
        # So we keep track of indices which are less than threshold and we keep filtering it away
        selected[(inds + i) + 1] = 0

    return keep

Is there a way I can use to compare the performance differences from the GPU version?
Thanks again for your time