| 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}.") | |