|
2 | 2 |
|
3 | 3 | #include <numeric> |
4 | 4 | #include <algorithm> |
| 5 | +#include <vector> |
5 | 6 |
|
6 | 7 | ONNXRuntime::ONNXRuntime(const std::string& model_path, const std::vector<std::string>& input_names) |
7 | 8 | : env_(new Ort::Env(OrtLoggingLevel::ORT_LOGGING_LEVEL_WARNING, "onnx_runtime")), |
@@ -58,16 +59,16 @@ ONNXRuntime::~ONNXRuntime() {} |
58 | 59 |
|
59 | 60 | template <typename T> |
60 | 61 | ONNXRuntime::Tensor<T> ONNXRuntime::run(Tensor<T>& input, |
61 | | - const Tensor<long>& input_shapes, |
| 62 | + const Tensor<int64_t>& input_shapes, |
62 | 63 | unsigned long long batch_size) const { |
63 | 64 | std::vector<Ort::Value> tensors_in; |
64 | 65 | for (const auto& name : input_node_strings_) { |
65 | 66 | auto input_pos = variablePos(name); |
66 | 67 | auto value = input.begin() + input_pos; |
67 | | - std::vector<int64_t> input_dims; |
| 68 | + std::vector< int64_t> input_dims; |
68 | 69 | if (input_shapes.empty()) { |
69 | | - input_dims = input_node_dims_.at(name); |
70 | | - input_dims[0] = batch_size; |
| 70 | + input_dims = input_node_dims_.at(name); |
| 71 | + input_dims[0] = (int64_t) batch_size; |
71 | 72 | } else { |
72 | 73 | input_dims = input_shapes[input_pos]; |
73 | 74 | // rely on the given input_shapes to set the batch size |
@@ -117,4 +118,4 @@ size_t ONNXRuntime::variablePos(const std::string& name) const { |
117 | 118 | return iter - input_names_.begin(); |
118 | 119 | } |
119 | 120 |
|
120 | | -template ONNXRuntime::Tensor<float> ONNXRuntime::run(Tensor<float>&, const Tensor<long>&, unsigned long long) const; |
| 121 | +template ONNXRuntime::Tensor<float> ONNXRuntime::run(Tensor<float>&, const Tensor<int64_t>&, unsigned long long) const; |
0 commit comments