Get CUDA GPU with most available memory

def get_free_gpu():
    max_free = 0
    max_idx = 0

    rows = (
        subprocess.check_output(
            ["nvidia-smi", "--format=csv", "--query-gpu=memory.free"]
        )
        .decode("utf-8")
        .split("\n")
    )
    for i, row in enumerate(rows[1:-1]):
        mb = float(row.rstrip(" [MiB]"))

        if mb > max_free:
            max_idx = i
            max_free = mb

    return max_idx