Optimizing Pytorch implementation of mix-up augmentation

So I have this code here for implementing mix-up augmentation. It's incredibly slow and I'm not sure how to make it faster. It seems like there are some operations that are unavoidable and just by nature slow like scaling images by the weight which is 0.5 then summing up each cell seems like a very slow and unavoidable operation. I'm applying this to Reinforcement Learning so I could be augmenting 64 million images, which is why I need it to be a lot faster.

Note: Here's the original author's implementation but I would assume it's equally as slow as it's essentially the same.

import torch
import utils
import os
import torch.nn.functional as F
import torchvision.transforms as TF
import torchvision.datasets as datasets

dataloader = None
data_iter = None

def _load_data(
    sub_path: str, batch_size: int = 256, image_size: int = 84, num_workers: int = 16
):
    global data_iter, dataloader
    for data_dir in utils.load_config("datasets"):
        if os.path.exists(data_dir):
            fp = os.path.join(data_dir, sub_path)
            if not os.path.exists(fp):
                print(f"Warning: path {fp} does not exist, falling back to {data_dir}")
            dataloader = torch.utils.data.DataLoader(
                datasets.ImageFolder(
                    fp,
                    TF.Compose(
                        [
                            TF.RandomResizedCrop(image_size),
                            TF.RandomHorizontalFlip(),
                            TF.ToTensor(),
                        ]
                    ),
                ),
                batch_size=batch_size,
                shuffle=True,
                num_workers=num_workers,
                pin_memory=True,
            )
            data_iter = iter(dataloader)
            break
    if data_iter is None:
        raise FileNotFoundError(
            "failed to find image data at any of the specified paths"
        )
    print("Loaded dataset from", data_dir)


def _load_places(batch_size=256, image_size=84, num_workers=16, use_val=False):
    partition = "val" if use_val else "train"
    sub_path = os.path.join("places365_standard", partition)
    print(f"Loading {partition} partition of places365_standard...")
    _load_data(
        sub_path=sub_path,
        batch_size=batch_size,
        image_size=image_size,
        num_workers=num_workers,
    )


def _load_coco(batch_size=256, image_size=84, num_workers=16, use_val=False):
    sub_path = "COCO"
    print(f"Loading COCO 2017 Val...")
    _load_data(
        sub_path=sub_path,
        batch_size=batch_size,
        image_size=image_size,
        num_workers=num_workers,
    )

def _get_data_batch(batch_size):
    global data_iter
    try:
        imgs, _ = next(data_iter)
        if imgs.size(0) < batch_size:
            data_iter = iter(dataloader)
            imgs, _ = next(data_iter)
    except StopIteration:
        data_iter = iter(dataloader)
        imgs, _ = next(data_iter)
    return imgs.cuda()

def load_dataloader(batch_size, image_size, dataset="coco"):
    if dataset == "places365_standard":
        if dataloader is None:
            _load_places(batch_size=batch_size, image_size=image_size)
    elif dataset == "coco":
        if dataloader is None:
            _load_coco(batch_size=batch_size, image_size=image_size)
    else:
        raise NotImplementedError(
            f'overlay has not been implemented for dataset "{dataset}"'
        )

def random_mixup(x, dataset="coco"):
    """Randomly overlay an image from Places or COCO"""
    global data_iter
    alpha = 0.5

    load_dataloader(batch_size=x.size(0), image_size=x.size(-1), dataset=dataset)

    imgs = _get_data_batch(batch_size=x.size(0)).repeat(1, x.size(1) // 3, 1, 1)

    return ((1 - alpha) * (x / 255.0) + (alpha) * imgs) * 255.0