Pytorch batch-wise Augmentation

I'm trying to fix the class imbalance of my dataset for a classification-task by adding augmented images. As the network didn't improve, I noticed, that I'm transforming the whole dataset while not keeping the original image. What is the best method to fix that?

My training function looks like this (excerpt):

def train(mu,lr,batch_size,n_epochs,k,model,use_gpu,size_image,seed,num_workers,root):
    set_seed(seed, use_gpu)
    train_loader, test_loader, dataset_attributes = get_data(size_image,root,batch_size, num_workers)
    
    criteria = CrossEntropyLoss()
    optimizer = SGD(model.parameters(), lr=lr, momentum=0.9, weight_decay=mu, nesterov=True)
    best_acc = 0.0

    for epoch in tqdm(range(n_epochs), desc='epoch', position=0):
        t = time.time()
        optimizer = update_optimizer(optimizer, lr_schedule=dataset_attributes['lr_schedule'], epoch=epoch)

        loss_epoch_train, f1_epoch_train, acc_epoch_train, topk_acc_epoch_train = train_epoch(model, optimizer, train_loader,
                                                                              criteria, loss_train, f1_train, acc_train,
                                                                              topk_acc_train, k,
                                                                              dataset_attributes['n_train'],
                                                                              use_gpu)
        if acc_epoch_test > best_acc:
            best_acc = acc_epoch_test
            save(model, optimizer, epoch, os.path.join(save_dir, 'weights_best_acc.tar'))

This in an excerpt of my get_data function:

def get_data(size_image,root,batch_size, num_workers):

    transform = transforms.Compose(
        [MaxCenterCrop(),
        transforms.Resize(size_image),
        transforms.ToTensor()])


    trainset = Plantnet(root, 'images_train', transform=transform)
    testset = Plantnet(root, 'images_test', transform=transform)
    train_class_to_num_instances = Counter(trainset.targets)
    test_class_to_num_instances = Counter(testset.targets)
    ...
    sampler = WeightedRandomSampler(torch.DoubleTensor(weights), int(num_samples))
    trainloader = torch.utils.data.DataLoader(trainset, batch_size=batch_size,
                                                sampler=sampler,
                                              shuffle=False, num_workers=num_workers)

    testloader = torch.utils.data.DataLoader(testset, batch_size=batch_size,
                                             shuffle=False, num_workers=num_workers)
    

    return trainloader, testloader, dataset_attributes

Now my idea for an easy fix would be to add a transformed dataset and concatenate it to the original one. But I think this idea would have a bad impact on performance and wouldn't really fix the problem of class imbalance.

I'm thinking that applying the tranformation on each batch would make the most sense. But how do I add this to my code?