int64_t channel = input_type.getDimSize(3);

  auto avg_pool_output_type = RankedTensorType::get(
      {batch, 1, 1, channel}, input_type.getElementType());
  auto avg_pool = rewriter.create<TFL::AveragePool2DOp>(
      mean_op.getLoc(), avg_pool_output_type, input,
      rewriter.getI32IntegerAttr(height), rewriter.getI32IntegerAttr(width),
      rewriter.getStringAttr("VALID"), rewriter.getI32IntegerAttr(1),