How can I concatenate Tensorflow Dataset columns?

Question

I have a Keras model that takes an input layer with shape (n, 288, 1), of which 288 is the number of features. I am using a TensorFlow dataset tf.data.experimental.make_batched_features_dataset and my input layer will be (n, 1, 1) which means it gives one feature to the model at a time. How can I make an input tensor with the shape of (n, 288, 1)? I mean how can I use all my features in one tensor? Here is my code for the model:

def _gzip_reader_fn(filenames):
"""Small utility returning a record reader that can read gzip'ed files."""
return tf.data.TFRecordDataset(filenames, compression_type='GZIP')


def _input_fn(file_pattern, tf_transform_output, batch_size):
"""Generates features and label for tuning/training.

Args:
    file_pattern: input tfrecord file pattern.
    tf_transform_output: A TFTransformOutput.
    batch_size: representing the number of consecutive elements of returned
    dataset to combine in a single batch

Returns:
    A dataset that contains (features, indices) tuple where features is a
     dictionary of Tensors, and indices is a single Tensor of label indices.
 """

transformed_feature_spec = (
    tf_transform_output.transformed_feature_spec().copy())

dataset = tf.data.experimental.make_batched_features_dataset(
    file_pattern=file_pattern,
    batch_size=batch_size,
    features=transformed_feature_spec,
    reader=_gzip_reader_fn,
    label_key=features.transformed_name(features.LABEL_KEY))

return dataset

def _build_keras_model(nb_classes=2, input_shape, learning_rate):

# Keras needs the feature definitions at compile time.
input_shape = (288,1)
input_layer = keras.layers.Input(input_shape)

padding = 'valid'
if input_shape[0] < 60:
    padding = 'same'

conv1 = keras.layers.Conv1D(filters=6, kernel_size=7, padding=padding, activation='sigmoid')(input_layer)
conv1 = keras.layers.AveragePooling1D(pool_size=3)(conv1)
conv2 = keras.layers.Conv1D(filters=12, kernel_size=7, padding=padding, activation='sigmoid')(conv1)
conv2 = keras.layers.AveragePooling1D(pool_size=3)(conv2)

flatten_layer = keras.layers.Flatten()(conv2)
output_layer = keras.layers.Dense(units=nb_classes, activation='sigmoid')(flatten_layer)

model = keras.models.Model(inputs=input_layer, outputs=output_layer)

optimizer = keras.optimizers.Adam(lr=learning_rate)

# Compile Keras model
model.compile(loss='mean_squared_error', optimizer=optimizer, metrics=['accuracy'])
model.summary(print_fn=logging.info)

return model

This is the error:

tensorflow:Model was constructed with shape (None, 288, 1) for input Tensor("input_1:0", shape=(None, 288, 1), dtype=float32), but it was called on an input with incompatible shape (128, 1, 1).