File size: 575 Bytes
d023e59
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
from optimum_benchmark.system_utils import get_gpus

if __name__=="__main__":
    # Get the names of all GPU devices
    gpu_name = get_gpus()

    # If there are several devices, keep the name of device 0
    if isinstance(gpu_name, list):
        gpu_name = gpu_name[0]

    # Raise an error if the device is not H100
    if "NVIDIA H100" in gpu_name:
        print("At least one NVIDIA H100 GPU has been detected, launching the benchmark...")
    else:
        raise RuntimeError(f"This Docker container should be executed on NVIDIA H100 GPUs only, detected {gpu_name}.")