From b8f5cabaddc78d2a09cb0b4f7e0c238e1fce2edf Mon Sep 17 00:00:00 2001 From: alexdrydew Date: Sun, 8 May 2022 21:21:22 +0300 Subject: [PATCH] support gpu training --- cuda_gpu_config.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cuda_gpu_config.py b/cuda_gpu_config.py index deeaefb..6c4551b 100644 --- a/cuda_gpu_config.py +++ b/cuda_gpu_config.py @@ -13,4 +13,4 @@ def setup_gpu(gpu_num=None): tf.config.experimental.set_memory_growth(gpu, True) logical_devices = tf.config.experimental.list_logical_devices('GPU') - assert len(logical_devices) > 0, "Not enough GPU hardware devices available" + assert len(logical_devices) > 0 or gpu_num == "", "Not enough GPU hardware devices available"