| 151 | } |
| 152 | |
| 153 | void BufferMap::SetFromTfLite(int tensor_index, const TfLiteTensor* tensor) { |
| 154 | tensorflow::TensorShape shape; |
| 155 | int num_dims = tensor->dims->size; |
| 156 | for (int i = 0; i < num_dims; ++i) { |
| 157 | shape.AddDim(tensor->dims->data[i]); |
| 158 | } |
| 159 | // TODO(ahentz): we assume this is a new tensor and allocate a new buffer |
| 160 | // for it. This is not always the best approach. For example, this might |
| 161 | // be a reallocation after resizing tensors. In that case it would be |
| 162 | // preferable to somehow reuse the buffer. |
| 163 | BaseTfLiteTensorBuffer* buf; |
| 164 | if (tensor->type == kTfLiteString) { |
| 165 | buf = new StringTfLiteTensorBuffer(tensor); |
| 166 | } else { |
| 167 | buf = new TfLiteTensorBuffer(tensor); |
| 168 | } |
| 169 | tensorflow::Tensor t = tensorflow::TensorCApi::MakeTensor( |
| 170 | GetTensorFlowDataType(tensor->type), shape, buf); |
| 171 | buf->Unref(); |
| 172 | |
| 173 | id_to_tensor_[tensor_index] = std::move(t); |
| 174 | owned_by_tf_.erase(tensor_index); |
| 175 | } |
| 176 | |
| 177 | void BufferMap::SetFromTensorFlow(int tensor_index, tensorflow::Tensor tensor) { |
| 178 | id_to_tensor_[tensor_index] = std::move(tensor); |