You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
18 lines
391 B
18 lines
391 B
5 months ago
|
#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
|