Training Model SSD300_VGG16 In Torchvision.Models

Hi,
I want to use the model SSD300_VGG16 from Torchvision.models.detection.ssd300_vgg16 with my custom dataset.
But when I trained this model, the loss didn’t decrease. And I checked inside the model.
I found out there are problems here.

Follow this code, torch.arange() is a problem because I have too many boxes with 0, and 1 box with another index. So location loss is not right.

Please check and help me. Thank you.