Decode output of nvmlDeviceGetName to avoid JSON serialize issue (#240

)
huggingface · Aug 15, 2024 · fd1d0f8 · fd1d0f8
1 parent cb30120
commit fd1d0f8
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/optimum_benchmark/system_utils.py b/optimum_benchmark/system_utils.py
@@ -95,7 +95,9 @@ def get_gpus():
         pynvml.nvmlInit()
         for i in range(pynvml.nvmlDeviceGetCount()):
             handle = pynvml.nvmlDeviceGetHandleByIndex(i)
-            gpus.append(pynvml.nvmlDeviceGetName(handle))
+            gpu = pynvml.nvmlDeviceGetName(handle)
+            gpu = gpu.decode("utf-8") if isinstance(gpu, bytes) else gpu # Older pynvml may return bytes 
+            gpus.append(gpu)
         pynvml.nvmlShutdown()
 
     elif is_rocm_system():