diff options
| author |  Yu-Cheng Ling <ycling@google.com> | 2018-02-16 15:16:38 -0800 |
|---|---|---|
| committer |  TensorFlower Gardener <gardener@tensorflow.org> | 2018-02-16 15:20:57 -0800 |
| commit | f4fd28d275887da0918f411dff2bced18a18c785 (patch) | |
| tree | 30d909f804653a697986379118f324e062776647 /tensorflow/contrib/lite/arena_planner.cc | |
| parent | 047c7cbc1c95f7423ee44c841eba6b0b3ec691bb (diff) | |
TFLite Conv2D: Create temporary tensors in Prepare phase.
PiperOrigin-RevId: 186053793
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()); |