aboutsummaryrefslogtreecommitdiffhomepage
path: root/tensorflow/contrib/lite/kernels/conv.cc
diff options
context:
space:
mode:
authorGravatar A. Unique TensorFlower <gardener@tensorflow.org>2018-05-17 11:47:16 -0700
committerGravatar TensorFlower Gardener <gardener@tensorflow.org>2018-05-17 11:50:17 -0700
commit622168b154be1aa60fe94fc840ad82803c29e5d3 (patch)
tree518100dddfea59de5711fc56a7e13bd3a3a488c7 /tensorflow/contrib/lite/kernels/conv.cc
parent637707c626404c0ab0b096fe7a57c5a2b9bc7719 (diff)
Support 1x1x1xN bias sizes in TFLite's convolution and FC layers.
PiperOrigin-RevId: 197027135
Diffstat (limited to 'tensorflow/contrib/lite/kernels/conv.cc')
-rw-r--r--tensorflow/contrib/lite/kernels/conv.cc3
1 files changed, 1 insertions, 2 deletions
diff --git a/tensorflow/contrib/lite/kernels/conv.cc b/tensorflow/contrib/lite/kernels/conv.cc
index 3b467b3aa2..2b7e455e3e 100644
--- a/tensorflow/contrib/lite/kernels/conv.cc
+++ b/tensorflow/contrib/lite/kernels/conv.cc
@@ -212,8 +212,7 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
} else {
TF_LITE_ENSURE_EQ(context, bias->type, data_type);
}
- TF_LITE_ENSURE_EQ(context, bias->dims->size, 1);
- TF_LITE_ENSURE_EQ(context, bias->dims->data[0], filter->dims->data[0]);
+ TF_LITE_ENSURE_EQ(context, NumElements(bias), SizeOfDimension(filter, 0));
}
int channels_out = filter->dims->data[0];