浏览代码

[bug-fix] Adjust MKL and KMP Settings (#4443)

* Adjust KMP and interop thread settings

* Don't set affinity, seems to slow down on low-CPU machines
/MLA-1734-demo-provider
GitHub 4 年前
当前提交
da3b0090
共有 1 个文件被更改,包括 4 次插入1 次删除
  1. 5
      ml-agents/mlagents/torch_utils/torch.py

5
ml-agents/mlagents/torch_utils/torch.py


import os
# Detect availability of torch package here.
# NOTE: this try/except is temporary until torch is required for ML-Agents.
try:

torch.set_num_threads(1)
torch.set_num_interop_threads(2)
os.environ["KMP_BLOCKTIME"] = "0"
# Known PyLint compatibility with PyTorch https://github.com/pytorch/pytorch/issues/701
# pylint: disable=E1101

正在加载...
取消
保存