Skip to content

Commit 10ebaab

Browse files
authored
Automated sync from github.com/tensorflow/tensorflow (#3031)
BUG=automated sync from upstream NO_CHECK_TFLITE_FILES=automated sync from upstream
1 parent 87d5152 commit 10ebaab

File tree

1 file changed

+5
-3
lines changed

1 file changed

+5
-3
lines changed

tensorflow/lite/kernels/internal/reference/integer_ops/fully_connected.h

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -42,12 +42,14 @@ void FullyConnectedPerChannel(
4242
const int32_t output_activation_min = params.quantized_activation_min;
4343
const int32_t output_activation_max = params.quantized_activation_max;
4444
TFLITE_DCHECK_GE(filter_shape.DimensionsCount(), 2);
45-
TFLITE_DCHECK_EQ(output_shape.DimensionsCount(), 2);
45+
TFLITE_DCHECK_GE(output_shape.DimensionsCount(), 1);
4646

4747
TFLITE_DCHECK_LE(output_activation_min, output_activation_max);
4848
const int filter_dim_count = filter_shape.DimensionsCount();
49-
const int batches = output_shape.Dims(0);
50-
const int output_depth = output_shape.Dims(1);
49+
50+
const int output_dim_count = output_shape.DimensionsCount();
51+
const int batches = FlatSizeSkipDim(output_shape, output_dim_count - 1);
52+
const int output_depth = output_shape.Dims(output_dim_count - 1);
5153
TFLITE_DCHECK_LE(output_depth, filter_shape.Dims(filter_dim_count - 2));
5254
const int accum_depth = filter_shape.Dims(filter_dim_count - 1);
5355
for (int b = 0; b < batches; ++b) {

0 commit comments

Comments
 (0)