pytorch

1
#include <gtest/gtest.h>
2

3
#include <torch/torch.h>
4

5
// NOTE: This test suite exists to make sure that common `torch::` functions
6
// can be used without additional includes beyond `torch/torch.h`.
7

8
TEST(TorchIncludeTest, GetSetNumThreads) {
9
  torch::init_num_threads();
10
  torch::set_num_threads(2);
11
  torch::set_num_interop_threads(2);
12
  torch::get_num_threads();
13
  torch::get_num_interop_threads();
14
}
15