WeightedRandomSampler with different target shape

Hello,

I am training an Object Detection Network. The dataset contains 2700 images and is highly imbalanced. the following dictionary indicates the class_id and the corresponding number of samples in the training dataset.

samples = {'1': 306, '2': 60, '3': 133, '4': 14197, '6': 3886, '7': 69}

Therefore, I want to use the WeightedRandomSampler. According to this post, I created the weights:

weights = torch.from_numpy(1. / np.asarray(list(samples.values())))

Then it is mentioned that these weights should be applied to the targets. However, I am not doing classification, but Object Detection using a modified CenterNet architecture. Therefore, for each image in the dataset, my target has a shape: [128, 384, 16]. This target contains where the number 16 comes from num_classes+4+1+4+1. Not all the images have all the classes.

Should I just loop through the dataset and create an array that contains just the class_labels and apply the weights to this array?
In this case, how do I use the WeighedRandomSampler?