mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 00:21:07 +01:00
25 lines
443 B
C++
25 lines
443 B
C++
#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
|