import numpy as np
import keras

from custom_layers import ScaleLayer  # Registers Guide>ScaleLayer for load_model().


sample = np.ones((2, 3), dtype="float32")
model = keras.saving.load_model("custom_scale_model.keras")
prediction = np.asarray(model(sample))

print(f"custom layer: {model.get_layer('scale_features').__class__.__name__}")
print(f"prediction shape: {prediction.shape}")
print(f"first prediction: {prediction[0].tolist()}")