Once you've created a tf.data.Dataset, you can compile and fit the model as before:

model.compile(optimizer=Adam(3e-5))  # No loss argument!