support gpu training

SiLiKhon · May 8, 2022 · b8f5cab · b8f5cab
1 parent 89d6d2e
commit b8f5cab
Showing 1 changed file with 1 addition and 1 deletion.
diff --git 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"