benchmarks/torch_ctl_benchmark/conv_model_benchmark.py/run_keras_custom_training_loop
def run_keras_custom_training_loop():
    keras_model = keras.Sequential(
        [
            layers.Input(shape=input_shape),
            layers.Conv2D(32, kernel_size=(3, 3), activation="relu"),
            layers.MaxPooling2D(pool_size=(2, 2)),
            layers.Flatten(),
            layers.Dense(num_classes),
            layers.Softmax(),
        ]
    )
    optimizer = optim.Adam(keras_model.parameters(), lr=0.001)
    loss_fn = nn.CrossEntropyLoss()
    train_loop(
        keras_model,
        train_loader,
        num_epochs=num_epochs,
        optimizer=optimizer,
        loss_fn=loss_fn,
        framework="keras",
    )