Cannot convert model with WeightNormalization from TFAddons: input resource[0] expected type resource != bool

I'm trying to convert to TFLite a model with weight normalization. However, an error is thrown during conversion with both TOCO or MLIR.

Model:

def conv_bn_relu(x, filters, kernel_size, strides, name, dilation_rate=(1, 1), use_bn=True, use_relu=True, use_wn=False, use_leakyrelu=False, padding = 'same'):
    #print('conv_bn_relu')
    if use_wn:
        x = WeightNormalization(Conv2D(filters=filters, 
                                       kernel_size=kernel_size, 
                                       strides=strides, 
                                       padding=padding, 
                                       dilation_rate=dilation_rate,
                                       name=name), data_init=False)(x)
    else:
        x = Conv2D(filters=filters, 
                   kernel_size=kernel_size, 
                   strides=strides, 
                   padding=padding, 
                   dilation_rate=dilation_rate,
                   name=name)(x)

    if use_bn:
        x = BatchNormalization(name='%s_bn' % (name))(x)        
    if use_relu:
        x = Activation('relu', name='%s_relu' % (name))(x)
    
    
    if use_leakyrelu:
        x = tf.keras.layers.LeakyReLU(alpha=0.2, name='%s_leakyrelu' % (name))(x)

    return x

Conversion snippet:

converter = tf.lite.TFLiteConverter.from_keras_model(generator_model)
tflite_model = converter.convert()
open(_MODEL_FILENAME_, "wb").write(tflite_model)

Error thrown:

~/anaconda3/envs/tf23fix/lib/python3.7/site-packages/tensorflow/lite/python/convert.py in toco_convert_protos(model_flags_str, toco_flags_str, input_data_str, debug_info_str, enable_mlir_converter)
    200       return model_str
    201     except Exception as e:
--> 202       raise ConverterError(str(e))
    203 
    204   if distutils.spawn.find_executable(_toco_from_proto_bin) is None:

ConverterError: input resource[0] expected type resource != bool, the type of functional_3_weight_normalization_cond_assert_equal_1_readvariableop_resource[0]
	In {{node functional_3/weight_normalization/cond/AssignVariableOp}}