import torch


def prepare_batch(batch, device, non_blocking, new_shape=None):
    """Prepare the batch data for training/inference, move data to GPU, reshape
    the target if necessary.

    Args:
        batch (torch.Tensor): A batch of data.
        device (torch.device or str): Device to load the backbone and data.
        non_blocking (bool): Whether tries to convert asynchronously with
            respect to the host if possible.
            https://pytorch.org/docs/stable/tensors.html#torch.Tensor.to
        new_shape (tuple): The new shape of the target variable, sometimes
            necessary for certain API calls.

    Returns:
        (torch.Tensor, torch.Tensor)

    """

    x = batch['image'].to(device, dtype=torch.float, non_blocking=non_blocking)

    y = batch['target'].to(device, dtype=torch.float, non_blocking=non_blocking)

    if new_shape:
        y = y.view(*new_shape)

    return x, y