diff options
| author |  Yu-Cheng Ling <ycling@google.com> | 2018-02-20 19:55:06 -0800 |
|---|---|---|
| committer |  TensorFlower Gardener <gardener@tensorflow.org> | 2018-02-20 19:59:16 -0800 |
| commit | 395616ff770318bfe19a9722bc8bc9d792779235 (patch) | |
| tree | 8d314dfaecdaee97fcaa9ecfafb150dfce5deffb /tensorflow/contrib/lite/arena_planner.cc | |
| parent | a75d7bf43bf1ff44566b7587e978a0c22b2ce171 (diff) | |
TFLite Conv2D: Create temporary tensors in Prepare phase.
PiperOrigin-RevId: 186402268
Diffstat (limited to 'tensorflow/contrib/lite/arena_planner.cc')
| -rw-r--r-- | tensorflow/contrib/lite/arena_planner.cc | 5 |
1 files changed, 5 insertions, 0 deletions
diff --git a/tensorflow/contrib/lite/arena_planner.cc b/tensorflow/contrib/lite/arena_planner.cc index 87b17c338e..8e47e2375e 100644 --- a/tensorflow/contrib/lite/arena_planner.cc +++ b/tensorflow/contrib/lite/arena_planner.cc @@ -128,6 +128,11 @@ TfLiteStatus ArenaPlanner::PlanAllocations() { } TfLiteStatus ArenaPlanner::ExecuteAllocations(int first_node, int last_node) { + // Grow the size of `allocs_` if necessary. This allows allocating temporary + // tensors in op's `prepare` function. + TF_LITE_ENSURE(context_, graph_info_->num_tensors() >= allocs_.size()); + allocs_.resize(graph_info_->num_tensors()); + TF_LITE_ENSURE_STATUS(CalculateAllocations(first_node, last_node)); TF_LITE_ENSURE_STATUS(Commit()); |