import torch


def pytorch_preprocess(batch):
    """
    The scaling procedure for all the pretrained models from torchvision is described in the docs
    https://pytorch.org/docs/stable/torchvision/models.html
    """
    batch = (batch + 1) * 0.5  # [-1, 1] -> [0, 1]

    batch_color_transformed = []
    batch = torch.stack(batch_color_transformed, 0)

    batch = torch.clamp(batch, 0, 1)
    mean = torch.tensor([.485, .456, .406], dtype=batch.dtype, device=batch.device)[None, :, None, None]
    batch = batch.sub(mean)  # subtract mean
    std = torch.tensor([.229, .224, .225], dtype=batch.dtype, device=batch.device)[None, :, None, None]
    batch = batch.div(std)
    return batch