commit 1f4a89b2d88f801f9d29342be1de8da97a755a64
author Marat Dukhan <maratek@google.com> Mon Mar 09 00:47:09 2020 -0700
committer TensorFlower Gardener <gardener@tensorflow.org> Mon Mar 09 00:50:19 2020 -0700
tree 71b20972ee0ee821a49631ba338946dfc0c43d8b
parent 839dae0b4657f660473689fb7bb127be5e58d197

Support non-unit batch with FC operator in GPU delegate

PiperOrigin-RevId: 299779651
Change-Id: I68da12096ea646144bc9cd70dbf4eca0aaa55c91

tensorflow/lite/delegates/gpu/common/model_builder.cc

diff --git a/tensorflow/lite/delegates/gpu/common/model_builder.cc b/tensorflow/lite/delegates/gpu/common/model_builder.cc
index d1521ff..b2bd6f0 100644
--- a/tensorflow/lite/delegates/gpu/common/model_builder.cc
+++ b/tensorflow/lite/delegates/gpu/common/model_builder.cc
@@ -1281,7 +1281,8 @@
       conv = graph->NewNode();  // reset conv pointer!
       Value<TensorRef<BHWC>>* reshaped_value = graph->NewValue();
       reshaped_value->tensor.type = DataType::FLOAT32;
-      reshaped_value->tensor.shape = BHWC(1, 1, 1, weights.shape.w);
+      reshaped_value->tensor.shape =
+          BHWC(input->tensor.shape.b, 1, 1, weights.shape.w);
       RETURN_IF_ERROR(graph->SetProducer(reshape->id, reshaped_value->id));
       reshape->operation.type = ToString(OperationType::RESHAPE);
       ReshapeAttributes attr;