File tree Expand file tree Collapse file tree 1 file changed +5
-3
lines changed
tensorflow/lite/kernels/internal/reference/integer_ops Expand file tree Collapse file tree 1 file changed +5
-3
lines changed Original file line number Diff line number Diff 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) {
You can’t perform that action at this time.
0 commit comments