frozenleaves/pytorch
1
0
Fork 0
mirror of https://github.com/pytorch/pytorch.git synced 2025-10-20 21:14:14 +08:00
Code Issues Packages Projects Releases Wiki Activity

Add get/set_num_interop_threads into torch.h include (#20659)

Browse Source 
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/20659
ghimport-source-id: 4858d03a9f89c613f64901c3430a7b212f76eb95

Reviewed By: dzhulgakov

Differential Revision: D15399780

Pulled By: ilia-cher

fbshipit-source-id: c1c3cb628c5ee664468f9d181bcd76a5105a89fd
This commit is contained in:
Ilia Cherniavskii
2019-05-20 00:31:39 -07:00
committed by Facebook Github Bot
parent 4598729399
commit 5835165ce3
2 changed files with 8 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
test/cpp/api/torch_include.cpp
View File

@ -9,4 +9,6 @@ TEST(TorchIncludeTest, GetSetNumThreads) {
torch::init_num_threads();
torch::set_num_threads(2);
ASSERT_EQ(torch::get_num_threads(), 2);
torch::set_num_interop_threads(2);
ASSERT_EQ(torch::get_num_interop_threads(), 2);
}

6
torch/csrc/api/include/torch/utils.h
View File

@ -26,4 +26,10 @@ using at::get_num_threads;
// Sets the number of threads to be used in parallel region.
using at::set_num_threads;
// Returns the number of threads used for inter-op parallelism.
using at::get_num_interop_threads;
// Sets the number of threads to be used for inter-op parallelism.
using at::set_num_interop_threads;
} // namespace torch