add strip for CUDA_VISIBLE_DEVICES

pull/740/head
huangyuxin 3 years ago
parent eacad8cf60
commit fccecf9976

@ -157,7 +157,7 @@ class Autolog:
model_precision="fp32"): model_precision="fp32"):
import auto_log import auto_log
pid = os.getpid() pid = os.getpid()
if (os.environ['CUDA_VISIBLE_DEVICES'] != ''): if (os.environ['CUDA_VISIBLE_DEVICES'].strip() != ''):
gpu_id = int(os.environ['CUDA_VISIBLE_DEVICES'].split(',')[0]) gpu_id = int(os.environ['CUDA_VISIBLE_DEVICES'].split(',')[0])
infer_config = inference.Config() infer_config = inference.Config()
infer_config.enable_use_gpu(100, gpu_id) infer_config.enable_use_gpu(100, gpu_id)

Loading…
Cancel
Save