1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
| #pragma once
|
| #include <ATen/Parallel.h>
| #include <c10/core/thread_pool.h>
|
| namespace at {
|
| class CAFFE2_API PTThreadPool : public c10::ThreadPool {
| public:
| explicit PTThreadPool(
| int pool_size,
| int numa_node_id = -1)
| : c10::ThreadPool(pool_size, numa_node_id, [](){
| c10::setThreadName("PTThreadPool");
| at::init_num_threads();
| }) {}
| };
|
| } // namespace at
|
|
