diff --git a/test/test_model_modules/test_flatten_tail.py b/test/test_model_modules/test_flatten_tail.py
index 09308ad44751d358e1b05fbf27f7fba0f4ce1818..83861be561fbe164d09048f1b748b51977b2fc27 100644
--- a/test/test_model_modules/test_flatten_tail.py
+++ b/test/test_model_modules/test_flatten_tail.py
@@ -1,7 +1,8 @@
+import tensorflow
 import tensorflow.keras as keras
 import pytest
 from mlair.model_modules.flatten import flatten_tail, get_activation
-
+from tensorflow.python.keras.layers.advanced_activations import ELU, ReLU
 
 class TestGetActivation:
 
@@ -18,10 +19,13 @@ class TestGetActivation:
     def test_sting_act_unknown(self, model_input):
         with pytest.raises(ValueError) as einfo:
             get_activation(model_input, activation='invalid_activation', name='String')
-        assert 'Unknown activation function:invalid_activation' in str(einfo.value)
+        assert 'Unknown activation function: invalid_activation. ' \
+               'Please ensure this object is passed to the `custom_objects` argument. ' \
+               'See https://www.tensorflow.org/guide/keras/save_and_serialize#registering_the_custom_object ' \
+               'for details.' in str(einfo.value)
 
     def test_layer_act(self, model_input):
-        x_in = get_activation(model_input, activation=keras.layers.advanced_activations.ELU, name='adv_layer')
+        x_in = get_activation(model_input, activation=ELU, name='adv_layer')
         act = x_in._keras_history[0]
         assert act.name == 'adv_layer'
 
@@ -44,7 +48,7 @@ class TestFlattenTail:
         return element
 
     def test_flatten_tail_no_bound_no_regul_no_drop(self, model_input):
-        tail = flatten_tail(input_x=model_input, inner_neurons=64, activation=keras.layers.advanced_activations.ELU,
+        tail = flatten_tail(input_x=model_input, inner_neurons=64, activation=ELU,
                             output_neurons=2, output_activation='linear',
                             reduction_filter=None,
                             name='Main_tail',
@@ -67,10 +71,10 @@ class TestFlattenTail:
         flatten = self.step_in(inner_dense)
         assert flatten.name == 'Main_tail'
         input_layer = self.step_in(flatten)
-        assert input_layer.input_shape == (None, 7, 1, 2)
+        assert input_layer.input_shape == [(None, 7, 1, 2)]
 
     def test_flatten_tail_all_settings(self, model_input):
-        tail = flatten_tail(input_x=model_input, inner_neurons=64, activation=keras.layers.advanced_activations.ELU,
+        tail = flatten_tail(input_x=model_input, inner_neurons=64, activation=ELU,
                             output_neurons=3, output_activation='linear',
                             reduction_filter=32,
                             name='Main_tail_all',
@@ -84,36 +88,40 @@ class TestFlattenTail:
         final_dense = self.step_in(final_act)
         assert final_dense.name == 'Main_tail_all_out_Dense'
         assert final_dense.units == 3
-        assert isinstance(final_dense.kernel_regularizer, keras.regularizers.L1L2)
+        assert isinstance(final_dense.kernel_regularizer, keras.regularizers.L2)
 
         final_dropout = self.step_in(final_dense)
         assert final_dropout.name == 'Main_tail_all_Dropout_2'
         assert final_dropout.rate == 0.35
 
         inner_act = self.step_in(final_dropout)
-        assert inner_act.get_config() == {'name': 'activation_1', 'trainable': True, 'activation': 'tanh'}
+        assert inner_act.get_config() == {'name': 'activation', 'trainable': True,
+                                          'dtype': 'float32', 'activation': 'tanh'}
 
         inner_dense = self.step_in(inner_act)
         assert inner_dense.units == 64
-        assert isinstance(inner_dense.kernel_regularizer, keras.regularizers.L1L2)
+        assert isinstance(inner_dense.kernel_regularizer, keras.regularizers.L2)
 
         inner_dropout = self.step_in(inner_dense)
-        assert inner_dropout.get_config() == {'name': 'Main_tail_all_Dropout_1', 'trainable': True, 'rate': 0.35,
+        assert inner_dropout.get_config() == {'name': 'Main_tail_all_Dropout_1', 'trainable': True,
+                                              'dtype': 'float32', 'rate': 0.35,
                                               'noise_shape': None, 'seed': None}
 
         flatten = self.step_in(inner_dropout)
-        assert flatten.get_config() == {'name': 'Main_tail_all', 'trainable': True, 'data_format': 'channels_last'}
+        assert flatten.get_config() == {'name': 'Main_tail_all', 'trainable': True,
+                                        'dtype': 'float32', 'data_format': 'channels_last'}
 
         reduc_act = self.step_in(flatten)
-        assert reduc_act.get_config() == {'name': 'Main_tail_all_conv_act', 'trainable': True, 'alpha': 1.0}
+        assert reduc_act.get_config() == {'name': 'Main_tail_all_conv_act', 'trainable': True,
+                                          'dtype': 'float32', 'alpha': 1.0}
 
         reduc_conv = self.step_in(reduc_act)
 
         assert reduc_conv.kernel_size == (1, 1)
         assert reduc_conv.name == 'Main_tail_all_Conv_1x1'
         assert reduc_conv.filters == 32
-        assert isinstance(reduc_conv.kernel_regularizer, keras.regularizers.L1L2)
+        assert isinstance(reduc_conv.kernel_regularizer, keras.regularizers.L2)
 
         input_layer = self.step_in(reduc_conv)
-        assert input_layer.input_shape == (None, 7, 1, 2)
+        assert input_layer.input_shape == [(None, 7, 1, 2)]