Function num_available_gpus

dali/test/python/test_utils_tensorflow.py:40–45

()

def num_available_gpus():
    local_devices = device_lib.list_local_devices()
    num_gpus = sum(1 for device in local_devices if device.device_type == "GPU")
    if not math.log2(num_gpus).is_integer():
        raise RuntimeError("Unsupported number of GPUs. This test can run on: 2^n GPUs.")
    return num_gpus


def available_gpus():