TF 2.4 issue for loading model with classes parameter

[irhallac]

I have used the example notebook to train a Linknet model. It is working with tensorflow 2.9.2. However I am having issues when I repeat the same steps with tensorflow 2.4.0. I checked the model.summary(). There is an issue in the softmax layer: Linknet(BACKBONE, encoder_weights='imagenet', classes=12, activation=activation)

TF 2.9.2:

conv2d (Conv2D)                (None, None, None,   1740        ['decoder_stage4c_relu[0][0]']   
                                12)                                                               

 softmax (Activation)           (None, None, None,   0           ['conv2d[0][0]']                 
                                12)                                                               

==================================================================================================
Total params: 11,523,285
Trainable params: 11,513,263
Non-trainable params: 10,022
__________________________________________________________________________

TF 2.4.0:

conv2d (Conv2D)                 (None, None, None, 1 1740        decoder_stage4c_relu[0][0]       
__________________________________________________________________________________________________
softmax (Activation)            (None, None, None, 1 0           conv2d[0][0]                     
==================================================================================================
Total params: 11,523,285
Trainable params: 11,513,263
Non-trainable params: 10,022
_________________________________________________________________