| 20 | |
| 21 | template<typename T> |
| 22 | Tensor BatchOneInput(const std::vector<Request>& requests, int index) { |
| 23 | TensorShape tensor_shape; |
| 24 | tensor_shape.AddDim(requests.size()); |
| 25 | |
| 26 | size_t single_input_size = 1; |
| 27 | auto input = requests[0].inputs[index].second; |
| 28 | for (int j = 0; j < input.dims(); ++j) { |
| 29 | tensor_shape.AddDim(input.dim_size(j)); |
| 30 | single_input_size *= input.dim_size(j); |
| 31 | } |
| 32 | |
| 33 | Tensor batched_tensor(input.dtype(), tensor_shape); |
| 34 | auto batched_tensor_flat = batched_tensor.flat<T>(); |
| 35 | size_t current_pos = 0; |
| 36 | for (auto req : requests) { |
| 37 | auto flat = req.inputs[index].second.flat<T>(); |
| 38 | memcpy(batched_tensor_flat.data() + current_pos, |
| 39 | flat.data(), single_input_size); |
| 40 | current_pos += single_input_size; |
| 41 | } |
| 42 | |
| 43 | return batched_tensor; |
| 44 | } |
| 45 | |
| 46 | template<typename T> |
| 47 | void SplitOneOutput(Response& batched_response, |