Test of a custom training loop

custom_train_test.py

import numpy as np
import tensorflow as tf
from tensorflow import keras

model = keras.models.Sequential([keras.layers.Dense(1, input_shape=[5])])
optimizer = keras.optimizers.SGD()

def step(model, optimizer, X_batch, y_batch):
    with tf.GradientTape() as tape:
        y_pred = model(X_batch)
    grads = tape.gradient(y_pred, model.trainable_weights)
    optimizer.apply_gradients(zip(grads, model.trainable_weights))

@tf.function
def train(model, optimizer, dataset):
    for X_batch, y_batch in dataset:
        step(model, optimizer, X_batch, y_batch)

X = np.random.rand(1000, 5)
y = np.random.rand(1000, 1)

dataset = tf.data.Dataset.from_tensor_slices((X, y)).batch(32).repeat(5)
train(model, optimizer, dataset)

Replacing the model with a subclassed model works just as well:

class MyModel(keras.models.Model):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.out = keras.layers.Dense(1)
    def call(self, inputs):
        return self.out(inputs)

model = MyModel()