pytorch/test/cpp/api/torch_include.cpp
Will Feng 0087069dce Use torch::get/set_num_threads without additional includes beyond torch/torch.h (#20176)
Summary:
Fixes https://github.com/pytorch/pytorch/issues/20130.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/20176

Differential Revision: D15275036

Pulled By: yf225

fbshipit-source-id: 0f04e1fbfed18c07030b20e92e957ef5f2b5707d
2019-05-09 06:08:27 -07:00

13 lines
342 B
C++

#include <gtest/gtest.h>
#include <torch/torch.h>
// NOTE: This test suite exists to make sure that common `torch::` functions
// can be used without additional includes beyond `torch/torch.h`.
TEST(TorchIncludeTest, GetSetNumThreads) {
torch::init_num_threads();
torch::set_num_threads(2);
ASSERT_EQ(torch::get_num_threads(), 2);
}