custom augmentation causes loss_rpn_loc to go to infinite

I have been trying to train faster RCNN with my own custom dataset and augmentation, however, the loss_rpn_loc keeps going to infinity.

I have already been able to train using my custom dataset by writing my own using the tutorial here:
https://pytorch.org/tutorials/beginner/basics/data_tutorial.html

I do not believe the dataset to be the issue - the training went fine when just using the custom dataset

However, when I add my own augmentation to shift the image around the loss_rpn_loc goes to infinity - rather unexpected since the plot doesn't go to infinity


the augmentation I have created will shift the image around randomly and update the bounding boxes accordingly, the code is included at the bottom of the post.

If I reduce the amount that the image could shift by, it works "better" - the loss won't become inf so soon.

The issue I take - which I have stated - is that the loss_rpn_loc will become inf rather abruptly and I am not sure how the augmentation would cause the error despite using it being the correlating factor

```

############################# XY shift Augmentation ############################

'''
info: applies the random shift in the x and y directions.

init:
inputs:
maxShift: tuple of the maximum shift (in pixels)

call:
inputs:
img: tensor to be augmented in the format [w,h,c] in 255
xyxy: bounding box coordinates for all objects

return:
img_out: the augmented image
xyxy: updated bounding box

'''

class ApplyRandomXYshift(object):
def __init__(self, maxShift = (-5,5)):

self.maxShift = maxShift

def __call__(self, img, xyxy):

img_height, img_width = img.shape[:2]

## fill will be zeros
img_out = torch.ones_like(img)*125

## x shift magnitude, y shift magnitude
x_shift, y_shift = tuple(torch.randint(self.maxShift[0], self.maxShift[1], (2,)))

## shift in x direction
xstart = max(0, x_shift)
xend = min(img_width, img_width + x_shift)

xstart1 = max(0, -x_shift)
xend1 = min(img_width, img_width - x_shift)

ystart = max(0, y_shift)
yend = min(img_height, img_height + y_shift)

ystart1 = max(0, -y_shift)
yend1 = min(img_height, img_height - y_shift)

img_out[ystart1:yend1,xstart1:xend1] = img[ystart:yend,xstart:xend]

## update box coordinates
xyxy[:, [0,2]] -= x_shift
xyxy[:, [1,3]] -= y_shift

xyxy[:, [0,2]] = np.clip(xyxy[:, [0,2]], 0, img_width)
xyxy[:, [1,3]] = np.clip(xyxy[:, [1,3]], 0, img_height)

return img_out, xyxy
```

"Instructions To Reproduce the Issue and Full Logs": not needed

"Your Environment": colab