def _float_tensor_to_list(output: FloatTensor) -> Union[
    List[List[float]], List[float]]:
    """
    Args:
        output (FloatTensor): _description_

    Returns:
        List[List[float]]: _description_
    """

    # Hardcoded to CPU always
    return output.detach().to("cpu").tolist()

    Args:
        output (FloatTensor): _description_

    Returns:
        List[List[float]]: _description_
    
src/marqo/s2_inference/s2_inference.py