未验证 提交 64c75771 编写于 作者: H Hui Zhang 提交者: GitHub

Merge pull request #740 from Jackwaterveg/fix_auto_log

fix the bug: can not use the CPU to test the model
......@@ -157,9 +157,13 @@ class Autolog:
model_precision="fp32"):
import auto_log
pid = os.getpid()
gpu_id = int(os.environ['CUDA_VISIBLE_DEVICES'].split(',')[0])
infer_config = inference.Config()
infer_config.enable_use_gpu(100, gpu_id)
if (os.environ['CUDA_VISIBLE_DEVICES'].strip() != ''):
gpu_id = int(os.environ['CUDA_VISIBLE_DEVICES'].split(',')[0])
infer_config = inference.Config()
infer_config.enable_use_gpu(100, gpu_id)
else:
gpu_id = None
infer_config = inference.Config()
autolog = auto_log.AutoLogger(
model_name=model_name,
model_precision=model_precision,
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册