ValueError: Found unexpected keys that do not correspond to any Model output

[callzhang]

System information

• OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Ubintu 18.04
• TensorFlow version and how it was installed (source or binary): 2.3.0
• TensorFlow-Addons version and how it was installed (source or binary): 0.12.0
• Python version: 3.8.5
• Is GPU used? (yes/no): both

Describe the bug

When using tfa.losses.SigmoidFocalCrossEntropy(from_logits=True, reduction='auto') the model raised ValueError:
ValueError: Found unexpected keys that do not correspond to any Model output: dict_keys(['class_name', 'config']). Expected: ['dense']

Code to reproduce the issue

...
loss = tfa.losses.SigmoidFocalCrossEntropy(from_logits=True, reduction='auto')
l_bert = bert.BertModelLayer.from_params(bert_params, name='albert' if albert else "bert")
l_input_ids = tf.keras.layers.Input(shape=(max_seq_length,), dtype='int32', name="input_ids")
l_mask_ids = tf.keras.layers.Input(shape=(max_seq_length,), dtype='int32', name="mask_ids")
l_token_type_ids = tf.keras.layers.Input(shape=(max_seq_length,), dtype='int32', name="token_type_ids")
output = l_bert([l_input_ids, l_token_type_ids], mask=l_mask_ids)
logits = tf.keras.layers.Dense(num_labels, kernel_initializer=tf.keras.initializers.TruncatedNormal(bert_params.initializer_range))(output)
model = tf.keras.Model(inputs=[l_input_ids, l_mask_ids, l_token_type_ids], outputs=logits)
model.build(input_shape=(None, max_seq_length))
model.fit(...)
...

Other info / logs