I am currently working on an implementation of the  Consistency regularization method proposed in the following paper.
[arXiv:2002.04724] Improved Consistency Regularization for GAN 1
Within this paper, when injecting real and generated images into the Discriminator, they apply data augmentation to each image and compute the loss value according to the following algorithm.
image1065×715 206 KB
If I only want to perform data augmentation on the real images, I can do so during batch data creation in torch.utils.data.Dataset, but how can I perform data augmentation on a Tensor generated from the Generator?
The following methods involve sending data between the CPU and the GPU, which causes some overhead and reduces computational speed.
def transforms_images(image_tensor):
    """transform tensor for consistency regularization
    :param image_tensor: torch.Tensor [B, C, H, W]
    """

    device = image_tensor.device
    _, _, H, W = image_tensor.shape

    cr_transforms = transforms.Compose([
        transforms.ToPILImage(),
        transforms.RandomResizedCrop((H, W)),
        transforms.ToTensor(),
        transforms.Normalize((0.5, ), (0.5, ))
    ])

    image_tensor = torch.cat([cr_transforms(image).unsqueeze_(0) 
                              for image in image_tensor.cpu()], dim=0)

    return image_tensor.to(device=device)

What’s the best way?