custom_objects.py

import tensorflow as tf
from tensorflow.keras import layers, models
from tensorflow.keras.saving import register_keras_serializable

@register_keras_serializable()
def trust_activation(x):
    return 0.5 + tf.sigmoid(x)

CUSTOM_OBJECTS = {
    'mse': tf.keras.losses.MeanSquaredError,
    'trust_activation': trust_activation
}