I’ve changed my implementation to this:
training_sets = data_augment(training_set)

# Storing the datasets before they are concatenated, 
# because we need the lengths of each
self.train_datasets = [MyDataset(
    features=self.features, 
    transform=transforms.Compose([ToTensor()])
) for _set in training_sets]
self.train_dataset = ConcatDataset(self.train_datasets)

# Creating weights for each concatenated dataset
weights = []
for d in self.train_datasets:
    dataset_len = len(d)
    weights.append(np.linspace(1, 3, dataset_len))
weights = np.concatenate(tuple(weights))
sampler = WeightedRandomSampler(weights, num_samples)
self.train_loader = DataLoader(self.train_dataset, sampler=sampler, num_workers=7, batch_size=self.batch_size)

I’ll let you know if it works better after my training run.
I am too lazy to debug this with synthetic data.