import torch


def _set_devices(model):
    """Set the device for the model.

    Args:
        model (model): The model to set up.

    Returns: The model on the correct device.

    """

    if torch.cuda.is_available():
        model = model.cuda()
        model = torch.nn.DataParallel(model)
    else:
        print("Using CPU")
        model = model.cpu()

    return model