The problem arises when the whole dataset is loaded by the GPU (instead of loading it with the CPU and only sending batches to the GPU).

I managed to fix this type of issue by adding with tf.device('cpu'): to the data loading process (just that, the remaining training should be done with the GPU).

Applied to your example code, something along the lines of:

model = build_model(
    filters=50,
    filter_step=1,
    stages=5,
    stage_steps=1,
    initial_convolutions=0,
    stacks=1,
)

print(model.summary()) 

with tf.device('cpu'):
    dataset = tf.data.Dataset.from_tensor_slices((X, y))
    dataset = dataset.batch(1)

model.fit(
    dataset,
    epochs=2**7,
    callbacks=[
        EarlyStopping(monitor="loss", patience=5, min_delta=1e-7, start_from_epoch=10),
        LearningRateScheduler(step_decay)
    ],
)

Thanks to @mrk's answer for pointing me in the right direction of the issue at hand. See also this question and the linked GitHub issue.