/usr/local/lib64/python3.6/site-packages/torch/include/caffe2/onnx
Edit: /usr/local/lib64/python3.6/site-packages/torch/include/caffe2/onnx/device.h (488B)
#pragma once
#include
#include
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 onnx
} // namespace caffe2
namespace std {
template <>
struct hash {
std::size_t operator()(const caffe2::onnx::DeviceType& k) const {
return std::hash()(static_cast(k));
}
};
} // namespace std