Custom Training Loop

import tensorflow as tf
import numpy as np

model = tf.keras.Sequential([tf.keras.layers.Dense(2, input_shape=(4,))])
optimizer = tf.keras.optimizers.Adam(1e-3)
loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)

train_ds = tf.data.Dataset.from_tensor_slices(
    (np.random.rand(200, 4).astype("float32"), np.random.randint(0, 2, (200,)))
).batch(32)

@tf.function
def train_step(x, y):
    with tf.GradientTape() as tape:
        logits = model(x, training=True)
        loss = loss_fn(y, logits)
    grads = tape.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(zip(grads, model.trainable_variables))
    return loss

for epoch in range(5):
    total = 0.0
    for x, y in train_ds:
        total += float(train_step(x, y))
    print(f"epoch {epoch} loss={total:.4f}")