refactor: on_gpu and to_device in baseclass

jina-ai · May 6, 2020 · ac4b10f · ac4b10f
1 parent d483605
commit ac4b10f
Showing 1 changed file with 1 addition and 4 deletions.
diff --git a/jina/executors/frameworks.py b/jina/executors/frameworks.py
@@ -20,10 +20,7 @@ def device(self):
         try:
             if self.framework == 'tensorflow':
                 import tensorflow as tf
-                if self.on_gpu:
-                    devices = tf.config.experimental.list_physical_devices(device_type='GPU')
-                else:
-                    devices = tf.config.experimental.list_physical_devices(device_type='CPU')
+                devices = tf.config.experimental.list_physical_devices(device_type='GPU' if self.on_gpu else 'CPU')
                 return devices[0]
             elif self.framework == 'paddlepaddle':
                 import paddle.fluid as fluid