def get_max_memory()

in src/optimum/nvidia/utils/nvml.py [0:0]

• 7 lines of code
• 2 McCabe index (conditional complexity)

def get_max_memory():
    fraction_device_map = {
        device_id: get_device_memory(device_id) * 0.7
        for device_id in range(get_device_count())
    }

    cpu_device_map = {"cpu": virtual_memory().available * 0.8}

    return fraction_device_map | cpu_device_map