custom_objects.py

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

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

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