import torch


def set_devices(model, loss, evaluator, device):
    """
    :param: model, loss, evaluation_loss : torch.nn.Module objects
    :param: device: torch.device,compatible string, or tuple/list of devices.
        If tuple, wrap the model using torch.nn.DataParallel and use the first
        specified device as the 'primary' one.
    :return: None

    Sets the model and loss the the specified device. Evaluation loss is performed on CPU.
    """
    print("Using device: ", device)
    if isinstance(device, (tuple, list)):
        model = torch.nn.DataParallel(model, device_ids=device)
        print("Using multi GPU compute on devices:", device)
        device = torch.device(device[0])
    device = torch.device(device)  # Will throw a more explicit error if the device specification is invalid

    model.to(device)
    loss.to(device)
    evaluator.loss.cpu()
    evaluator.model_device = device
    evaluator.model = model

    return model, evaluator