import torch


def format_output(obj, cuda_devices):
    """
    Format output based on whether model is using DataParallel or not.
    DataParallel necessitates objects to be gathered into GPU:0 to have
    dimension 0.
    This function will be used for scalar outputs of model forwards
    such as loss and n_exs.
    """
    if isinstance(cuda_devices, list):
        if not isinstance(obj, torch.Tensor):
            obj = torch.tensor(obj).cuda()
        return obj.unsqueeze(0)
    else:
        return obj