fix(helper): fix gpuutil exception

gnes-ai · Oct 15, 2019 · af7c885 · af7c885
1 parent b6837b1
commit af7c885
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 6 deletions.
diff --git a/gnes/encoder/image/inception.py b/gnes/encoder/image/inception.py
@@ -35,11 +35,12 @@ def __init__(self, model_dir: str,
         self.inception_size_y = 299
 
     def post_init(self):
+        import os
+        os.environ['CUDA_VISIBLE_DEVICES'] = str(get_first_available_gpu())
+
         import tensorflow as tf
         from .inception_cores.inception_v4 import inception_v4
         from .inception_cores.inception_utils import inception_arg_scope
-        import os
-        os.environ['CUDA_VISIBLE_DEVICES'] = str(get_first_available_gpu())
         g = tf.Graph()
         with g.as_default():
             arg_scope = inception_arg_scope()

diff --git a/gnes/helper.py b/gnes/helper.py
@@ -18,8 +18,8 @@
 import logging
 import os
 import sys
-import time
 import threading
+import time
 from copy import copy
 from functools import wraps
 from itertools import islice
@@ -104,9 +104,7 @@ def get_first_available_gpu():
         if r:
             return r[0]
         raise ValueError
-    except ImportError:
-        return -1
-    except ValueError:
+    except Exception:
         return -1