| 12345678910111213141516171819202122 |
- #if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
- #pragma once
- #include <ATen/Parallel.h>
- #include <c10/core/thread_pool.h>
- namespace at {
- class TORCH_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
- #else
- #error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
- #endif // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
|
