fix single gpu training

RocketFlash · Jul 16, 2020 · 8b99e27 · 8b99e27
1 parent 4088aea
commit 8b99e27
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/tools/train.py b/tools/train.py
@@ -127,6 +127,8 @@ def main():
             gpu_ids_list = gpu_ids.split(',')
             n_gpu = len(gpu_ids_list)
         else:
+            os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" 
+            os.environ["CUDA_VISIBLE_DEVICES"] = '0'
             n_gpu = 1
             print('Use single gpu mode')