Fixes #14199: catch subprocess.CalledProcessError in get_gpus()

apache · Feb 21, 2019 · 8e29e18 · 8e29e18
1 parent fdbc433
commit 8e29e18
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 3 deletions.
diff --git a/benchmark/python/control_flow/rnn.py b/benchmark/python/control_flow/rnn.py
@@ -81,7 +81,7 @@ def _array(shape, ctx):
 def _get_gpus():
     try:
         re = subprocess.check_output(["nvidia-smi", "-L"], universal_newlines=True)
-    except OSError:
+    except (subprocess.CalledProcessError, OSError):
         return []
     return range(len([i for i in re.split('\n') if 'GPU' in i]))
 

diff --git a/example/image-classification/common/util.py b/example/image-classification/common/util.py
@@ -51,6 +51,6 @@ def get_gpus():
     """
     try:
         re = subprocess.check_output(["nvidia-smi", "-L"], universal_newlines=True)
-    except OSError:
+    except (subprocess.CalledProcessError, OSError):
         return []
     return range(len([i for i in re.split('\n') if 'GPU' in i]))
diff --git a/tools/bandwidth/test_measure.py b/tools/bandwidth/test_measure.py
@@ -24,7 +24,7 @@
 def get_gpus():
     try:
         re = subprocess.check_output(["nvidia-smi", "-L"], universal_newlines=True)
-    except OSError:
+    except (subprocess.CalledProcessError, OSError):
         return ''
     gpus = [i for i in re.split('\n') if 'GPU' in i]
     return ','.join([str(i) for i in range(len(gpus))])