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(
42
42
const int32_t output_activation_min = params.quantized_activation_min ;
43
43
const int32_t output_activation_max = params.quantized_activation_max ;
44
44
TFLITE_DCHECK_GE (filter_shape.DimensionsCount (), 2 );
45
- TFLITE_DCHECK_EQ (output_shape.DimensionsCount (), 2 );
45
+ TFLITE_DCHECK_GE (output_shape.DimensionsCount (), 1 );
46
46
47
47
TFLITE_DCHECK_LE (output_activation_min, output_activation_max);
48
48
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 );
51
53
TFLITE_DCHECK_LE (output_depth, filter_shape.Dims (filter_dim_count - 2 ));
52
54
const int accum_depth = filter_shape.Dims (filter_dim_count - 1 );
53
55
for (int b = 0 ; b < batches; ++b) {
You can’t perform that action at this time.
0 commit comments