import torch


def model_to(on_gpu, model):
    """Transfers model to cpu/gpu.

    Args:
        on_gpu (bool): Transfers model to gpu if True otherwise to cpu
        model (torch.nn.Module): PyTorch defined model.

    Returns:
        torch.nn.Module:
            The model after being moved to cpu/gpu.

    """
    if on_gpu:  # DataParallel work only for cuda
        model = torch.nn.DataParallel(model)
        return model.to("cuda")

    return model.to("cpu")