Simple Model with make_csv_data

I have the following CSV dataset (note that feature1 and feature3 are integer, and feature2 is a float, also note that I’d like to use make_csv_dataset because my input actually have 400 features and millions of rows):

feature1,feature2,feature3,label
0,1.3,2,0
1,2.7,5,1
1,2.7,5,1
...

I’m writing the following keras model and surprised that it doesn’t work:

_TRAINING_DATA_PATH = "data.csv"
_BATCH_SIZE = 64

ds = tf.data.experimental.make_csv_dataset(
    _TRAINING_DATA_PATH, batch_size=_BATCH_SIZE,
    select_columns = ["feature1","feature2","feature3","label"],
    label_name="label")

model = tf.keras.Sequential([
    tf.keras.Input(shape=(3, ), batch_size=_BATCH_SIZE),
    tf.keras.layers.Dense(16, activation='relu'),
    tf.keras.layers.Dense(8, activation='relu'),
    tf.keras.layers.Dense(1)])

model.compile(loss=tf.keras.losses.BinaryCrossentropy(from_logits=False),
              optimizer='adam',
              metrics=tf.metrics.BinaryAccuracy(threshold=0.0))

model.fit(ds, epochs=10)

Any idea why this code receives the error below? I saw some similar posts that suggests mapping the dataset to a different format but I feel like those kind of manipulations should not be necessary (also they don’t work with mixed typed features like ints and floats together). Curious how would one fix the code above?

ValueError                                Traceback (most recent call last)
<ipython-input-16-9f049ca8f25a> in <module>
----> 1 model.fit(ds, epochs=10)

1 frames
/usr/local/lib/python3.8/dist-packages/keras/engine/training.py in tf__train_function(iterator)
     13                 try:
     14                     do_return = True
---> 15                     retval_ = ag__.converted_call(ag__.ld(step_function), (ag__.ld(self), ag__.ld(iterator)), None, fscope)
     16                 except:
     17                     do_return = False

ValueError: in user code:

    File "/usr/local/lib/python3.8/dist-packages/keras/engine/training.py", line 1051, in train_function  *
        return step_function(self, iterator)
    File "/usr/local/lib/python3.8/dist-packages/keras/engine/training.py", line 1040, in step_function  **
        outputs = model.distribute_strategy.run(run_step, args=(data,))
    File "/usr/local/lib/python3.8/dist-packages/keras/engine/training.py", line 1030, in run_step  **
        outputs = model.train_step(data)
    File "/usr/local/lib/python3.8/dist-packages/keras/engine/training.py", line 889, in train_step
        y_pred = self(x, training=True)
    File "/usr/local/lib/python3.8/dist-packages/keras/utils/traceback_utils.py", line 67, in error_handler
        raise e.with_traceback(filtered_tb) from None
    File "/usr/local/lib/python3.8/dist-packages/keras/engine/input_spec.py", line 183, in assert_input_compatibility
        raise ValueError(f'Missing data for input "{name}". '

    ValueError: Missing data for input "input_5". You passed a data dictionary with keys ['feature1', 'feature2', 'feature3']. Expected the following keys: ['input_5']