Sequential Model Save/Load Problems

I have a Tensorflow 2.x model which is using the TF preprocessing layer (tf.keras.layers.DenseFeatures) and the distributional layer from TF probability (DistributionLambda)

def regression_deep1_proba2(preprocessing_layer, feature_layer_inputs, model_name='test_model'):


    model = tf.keras.Sequential([
        preprocessing_layer,
        tf.keras.layers.Dense(100, activation='relu', name='hidden_1'),
        tf.keras.layers.Dense(50, activation='relu', name='hidden_2'),
        tf.keras.layers.Dense(1 + 1, name='output'),
        tfp.layers.DistributionLambda(
            lambda t: tfd.LogNormal(loc=t[..., :1], scale=tf.math.softplus(0.05 * t[..., 1:]))
        ),
    ])

    # ____________________ COMPILE WITH  ____________________________________________
    optimizer = tf.keras.optimizers.Adam()
    negloglik = lambda y, p_y: -p_y.log_prob(y)

    metrics = [
        tf.keras.metrics.MeanAbsolutePercentageError()
        ]

    model.compile(
        loss=negloglik,
        optimizer=optimizer,
        metrics=metrics
    )

    # ____________________ CALLBACKS DEFINITION ___________________________________________
    tbCallBack = tf.keras.callbacks.TensorBoard(
        log_dir=f'./logs_regression/{model_name}',
        update_freq='batch',
        histogram_freq=1,
        embeddings_freq=1,
        write_graph=True,
        write_images=True
    )

    # Create a callback that saves the model's weights every 5 epochs
    cp_callback = tf.keras.callbacks.ModelCheckpoint(
        filepath=f'./weights.{model_name}.hdf5',
        verbose=1,
        save_weights_only=True,
        save_best_onlt=True,
        monitor='MeanSquaredError'
    )
    early_stop = tf.keras.callbacks.EarlyStopping(
        monitor='MeanSquaredError',
        patience=2
    )
    callbacks_list = [tbCallBack, cp_callback, early_stop]

    return model, callbacks_list

I can get some nice results for the regression problem with this model, but when I save it for further use I can't load it back anymore (I have tried all online tutorials and solutions, but nothing is working)!!

I can save it to a file (h5, tf, json etc...)
i.e.:

tf.keras.models.save_model(model, 'model_name.h5')

but when loading I get:

ValueError: ('We expected a dictionary here. Instead we got: ', <tf.Tensor 'Placeholder:0' shape=(None,) dtype=float32>)

I can't figure out what am I doing wrong - any help would be appreciated!

Also, I have tried all possible save extensions and backends: h5, tf, json, simple weights and other formats but none of them works ... I have even tried to do it on different systems: Mac, Ubuntu and on different Tensorflow versions: 2 and 2.1 ...

Of course, all the saving and loading works great for other models I use without the TF Probability layer (even the ones with a DenseFeatures layer).

Start with the regression_deep1_proba2 model construction and the tf.keras.models.save_model call shown in the issue, then reproduce loading the saved Sequential model containing DenseFeatures and DistributionLambda. Compare the save and load behavior across the mentioned formats and TensorFlow versions; done means the saved model can be loaded successfully for further use.