fix the bug: can not use the CPU to test the model

3 years ago · eacad8cf60
parent 8432b9b256
commit eacad8cf60
1 changed files with 7 additions and 3 deletions
--- a/deepspeech/utils/log.py
+++ b/deepspeech/utils/log.py
@ -157,9 +157,13 @@ class Autolog:
                 model_precision="fp32"):
        import auto_log
        pid = os.getpid()
-        gpu_id = int(os.environ['CUDA_VISIBLE_DEVICES'].split(',')[0])
-        infer_config = inference.Config()
-        infer_config.enable_use_gpu(100, gpu_id)
+        if (os.environ['CUDA_VISIBLE_DEVICES'] != ''):
+            gpu_id = int(os.environ['CUDA_VISIBLE_DEVICES'].split(',')[0])
+            infer_config = inference.Config()
+            infer_config.enable_use_gpu(100, gpu_id)
+        else:
+            gpu_id = None
+            infer_config = inference.Config()
        autolog = auto_log.AutoLogger(
            model_name=model_name,
            model_precision=model_precision,