|
|
|
@@ -61,8 +61,8 @@ def device_placement(framework, device_name='gpu:0'): |
|
|
|
if framework == Frameworks.tf: |
|
|
|
import tensorflow as tf |
|
|
|
if device_type == Devices.gpu and not tf.test.is_gpu_available(): |
|
|
|
logger.warning( |
|
|
|
'tensorflow cuda is not available, using cpu instead.') |
|
|
|
logger.debug( |
|
|
|
'tensorflow: cuda is not available, using cpu instead.') |
|
|
|
device_type = Devices.cpu |
|
|
|
if device_type == Devices.cpu: |
|
|
|
with tf.device('/CPU:0'): |
|
|
|
@@ -78,7 +78,8 @@ def device_placement(framework, device_name='gpu:0'): |
|
|
|
if torch.cuda.is_available(): |
|
|
|
torch.cuda.set_device(f'cuda:{device_id}') |
|
|
|
else: |
|
|
|
logger.warning('cuda is not available, using cpu instead.') |
|
|
|
logger.debug( |
|
|
|
'pytorch: cuda is not available, using cpu instead.') |
|
|
|
yield |
|
|
|
else: |
|
|
|
yield |
|
|
|
@@ -96,9 +97,7 @@ def create_device(device_name): |
|
|
|
if device_type == Devices.gpu: |
|
|
|
use_cuda = True |
|
|
|
if not torch.cuda.is_available(): |
|
|
|
logger.warning( |
|
|
|
'cuda is not available, create gpu device failed, using cpu instead.' |
|
|
|
) |
|
|
|
logger.info('cuda is not available, using cpu instead.') |
|
|
|
use_cuda = False |
|
|
|
|
|
|
|
if use_cuda: |
|
|
|
|