1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
| #pragma once
|
| #include <functional>
| #include <string>
|
| namespace caffe2 { namespace onnx {
|
| enum class DeviceType {CPU=0, CUDA=1};
|
| struct Device {
| Device(const std::string& spec);
| DeviceType type;
| int device_id{-1};
| };
|
| }}
|
| namespace std {
| template <> struct hash<caffe2::onnx::DeviceType> {
| std::size_t operator()(const caffe2::onnx::DeviceType &k) const {
| return std::hash<int>()(static_cast<int>(k));
| }
| };
| } // namespace std
|
|
