" Depth of input is not a multiple of input depth of filter" error when quantizing loaded model

I want to perform quantization-aware training on a keras classification model. The first training steps are performed without quantization and the model is saved. Then, the model is quantized and the training goes on.
Now, when I call quantize_model on the loaded model I get the following error:

Traceback (most recent call last):
File "train.py", line 121, in
model = tfmot.quantization.keras.quantize_model(model)
File "/usr/local/lib/python3.6/dist-packages/tensorflow_model_optimization/python/core/quantization/keras/quantize.py", line 138, in quantize_model
return quantize_apply(annotated_model)
File "/usr/local/lib/python3.6/dist-packages/tensorflow_model_optimization/python/core/quantization/keras/quantize.py", line 410, in quantize_apply
unwrapped_model, layer_quantize_map)
File "/usr/local/lib/python3.6/dist-packages/tensorflow_model_optimization/python/core/quantization/keras/default_8bit/default_8bit_quantize_layout_transform.py", line 73, in apply
set(layer_quantize_map.keys()), layer_quantize_map).transform()
File "/usr/local/lib/python3.6/dist-packages/tensorflow_model_optimization/python/core/quantization/keras/graph_transformations/model_transformer.py", line 551, in transform
transformed_model = keras.Model.from_config(self._config, custom_objects)
File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/training.py", line 2262, in from_config
config, custom_objects=custom_objects)
File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/functional.py", line 669, in from_config
config, custom_objects)
File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/functional.py", line 1285, in reconstruct_from_config
process_node(layer, node_data)
File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/functional.py", line 1233, in process_node
output_tensors = layer(input_tensors, **kwargs)
File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/base_layer.py", line 952, in __call__
input_list)
File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/base_layer.py", line 1091, in _functional_construction_call
inputs, input_masks, args, kwargs)
File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/base_layer.py", line 822, in _keras_tensor_symbolic_call
return self._infer_output_signature(inputs, args, kwargs, input_masks)
File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/base_layer.py", line 863, in _infer_output_signature
outputs = call_fn(inputs, *args, **kwargs)
File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/layers/convolutional.py", line 248, in call
outputs = self._convolution_op(inputs, self.kernel)
File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/util/dispatch.py", line 201, in wrapper
return target(*args, **kwargs)
File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/ops/nn_ops.py", line 1020, in convolution_v2
name=name)
File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/ops/nn_ops.py", line 1150, in convolution_internal
name=name)
File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/ops/nn_ops.py", line 2604, in _conv2d_expanded_batch
name=name)
File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/ops/gen_nn_ops.py", line 973, in conv2d
data_format=data_format, dilations=dilations, name=name)
File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/framework/op_def_library.py", line 750, in _apply_op_helper
attrs=attr_protos, op_def=op_def)
File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/framework/func_graph.py", line 592, in _create_op_internal
compute_device)
File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/framework/ops.py", line 3536, in _create_op_internal
op_def=op_def)
File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/framework/ops.py", line 2016, in __init__
control_input_ops, op_def)
File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/framework/ops.py", line 1856, in _create_c_op
raise ValueError(str(e))
ValueError: Depth of input (48) is not a multiple of input depth of filter (32) for '{{node conv2d_1/Conv2D}} = Conv2D[T=DT_FLOAT, data_format="NHWC", dilations=[1, 1, 1, 1], explicit_paddings=[], padding="VALID", strides=[1, 1, 1, 1], use_cudnn_on_gpu=true](Placeholder, conv2d_1/Conv2D/ReadVariableOp)' with input shapes: [?,35,35,48], [1,1,32,48].

Please note that the following error is not thrown if the quantization is performed on the model at the beginning (instead of the loaded model), even if the architecture is the same.
Some information:
tensorflow-gpu 2.4.0
tfmot v 0.5.0

model:

def inception_block(x, a, b, c, d, e, f):
    c1 = layers.Conv2D(a, 1, activation='relu')(x)
    c3r = layers.Conv2D(b, 1, activation='relu')(x)
    c5r = layers.Conv2D(c, 1, activation='relu')(x)
    pool = layers.MaxPool2D(3, 1, padding='same')(x)
    c3 = layers.Conv2D(d, 3, padding='same', activation='relu')(c3r)
    c5 = layers.Conv2D(e, 5, padding='same', activation='relu')(c5r)
    pool_proj = layers.Conv2D(f, 1, padding='same')(pool)
    return layers.Concatenate()([c1, c3, c5, pool_proj])

def build_model(N, input_size):
    if output_bias is not None:
        output_bias = tf.keras.initializers.Constant(output_bias)

    inputs = tf.keras.Input(shape=(input_size, input_size, N))

    x = layers.SeparableConv2D(32, 5, 2)(inputs)
    x = layers.BatchNormalization()(x)
    x = layers.ReLU(6.)(x)
    x = layers.MaxPool2D(3, 2)(x) 
    x = layers.Dropout(0.1)(x)   
    x = layers.SeparableConv2D(64, 3)(x)
    x = layers.BatchNormalization()(x)
    x = layers.ReLU(6)(x)
    x = layers.MaxPool2D(3, 2)(x)
    x = layers.Dropout(0.1)(x)
    x = inception_block(x, 32, 48, 32, 64, 64, 64)
    x = inception_block(x, 48, 48, 48, 64, 32, 24)
    x = layers.AveragePooling2D(5, 1)(x)
    x = layers.Dropout(0.5)(x)
    x = layers.Flatten()(x)
    x = layers.Dense(1, activation='sigmoid')(x)

    model = tf.keras.Model(inputs=inputs, outputs=x)

    return model

Training code for loaded model (throws error):

model = tf.keras.models.load_model(load_ckpt_path)
model = tfmot.quantization.keras.quantize_model(model)

Training code to quantize from the beginning (does not throw error):

model = build_model(input_size, N)
model = tfmot.quantization.keras.quantize_model(model)

EDIT: apparently, replacing the SeparableConv2D layers with Conv2D fixes the problem. What is the reason?