Keras + Tensorflow Model Optimization: TypeError: clone_model() got an unexpected keyword argument 'clone_function'

I'm trying Tensorflow Model Optimization in order to prune a simple Neural Network. Here's my code:

from __future__ import absolute_import, division, print_function, unicode_literals, unicode_literals
import tensorflow as tf
from tensorflow import keras
import numpy as np
import matplotlib.pyplot as plt

fashion_mnist = keras.datasets.fashion_mnist

(train_images, train_labels), (test_images, test_labels) = fashion_mnist.load_data()

class_names = ['T-shirt/top', 'Trouser', 'Pullover', 'Dress', 'Coat',
               'Sandal', 'Shirt', 'Sneaker', 'Bag', 'Ankle boot']

train_images = train_images / 255.0
test_images = test_images / 255.0

model = keras.Sequential([
    keras.layers.Flatten(input_shape=(28, 28)),
    keras.layers.Dense(128, activation=tf.nn.relu),
    keras.layers.Dense(10, activation=tf.nn.softmax)
])

import tensorflow_model_optimization as tfmot


pruning_schedule = tfmot.sparsity.keras.PolynomialDecay(
                        initial_sparsity=0.0, final_sparsity=0.5,
                        begin_step=2000, end_step=4000)

model_for_pruning = tfmot.sparsity.keras.prune_low_magnitude(model, pruning_schedule=pruning_schedule)


model_for_pruning.compile(optimizer='adam',
          loss='sparse_categorical_crossentropy',
          metrics=['accuracy'])

from tensorflow.keras.callbacks import TensorBoard


tensorboard=TensorBoard(log_dir='D:\Python\logs', histogram_freq=0,  
          write_graph=True, write_images=True)

model_for_pruning.fit(train_images, train_labels, epochs=5,callbacks=tensorboard)


#tensorboard --logdir D:\Python\logs 

I'm getting the following error:

File "<ipython-input-1-8f75575649d2>", line 52, in <module>
    model_for_pruning = tfmot.sparsity.keras.prune_low_magnitude(model, pruning_schedule=pruning_schedule)

  File "C:\Users\Rubens\Anaconda3\lib\site-packages\tensorflow_model_optimization\python\core\sparsity\keras\prune.py", line 152, in prune_low_magnitude
    to_prune, input_tensors=None, clone_function=_add_pruning_wrapper)

TypeError: clone_model() got an unexpected keyword argument 'clone_function'

That is, clone function does not belong to Keras' file models.py. I tried to add **kwargs to it, without success:

def clone_model(model, input_tensors=None,**kwargs):
"""Clone any `Model` instance.

Model cloning is similar to calling a model on new inputs,
except that it creates new layers (and thus new weights) instead
of sharing the weights of the existing layers.

# Arguments
    model: Instance of `Model`
        (could be a functional model or a Sequential model).
    input_tensors: optional list of input tensors
        to build the model upon. If not provided,
        placeholders will be created.

# Returns
    An instance of `Model` reproducing the behavior
    of the original model, on top of new inputs tensors,
    using newly instantiated weights.

# Raises
    ValueError: in case of invalid `model` argument value.
"""
if isinstance(model, Sequential):
    return _clone_sequential_model(model, input_tensors=input_tensors)
else:
    return _clone_functional_model(model, input_tensors=input_tensors)

This is the end of file prune.py, belonging to Tensorflow Model Optimization (notice clone_function=_strip_pruning_wrapper):

  def _strip_pruning_wrapper(layer):
    if isinstance(layer, pruning_wrapper.PruneLowMagnitude):
      # The _batch_input_shape attribute in the first layer makes a Sequential
      # model to be built. This makes sure that when we remove the wrapper from
      # the first layer the model's built state preserves.
      if not hasattr(layer.layer, '_batch_input_shape') and hasattr(
          layer, '_batch_input_shape'):
        layer.layer._batch_input_shape = layer._batch_input_shape
      return layer.layer
    return layer

  return keras.models.clone_model(
      model, input_tensors=None, clone_function=_strip_pruning_wrapper)

All libraries included are up-to-date. Any ideas on how to overcome this error ?

Thanks in advance

I found the answer. There is a tricky workaround: besides fixing the code to:

from tensorflow_model_optimization.sparsity import keras as sparsity

pruning_params = {
      'pruning_schedule': sparsity.PolynomialDecay(initial_sparsity=0.50,
                                                   final_sparsity=0.90,
                                                   begin_step=3,
                                                   end_step=end_step,
                                                   frequency=100)
}

pruned_model = tf.keras.Sequential([
    sparsity.prune_low_magnitude(
        l.Conv2D(32, 5, padding='same', activation='relu'),
        input_shape=input_shape,
        **pruning_params),
    l.MaxPooling2D((2, 2), (2, 2), padding='same'),
    l.BatchNormalization(),
    sparsity.prune_low_magnitude(
        l.Conv2D(64, 5, padding='same', activation='relu'), **pruning_params),
    l.MaxPooling2D((2, 2), (2, 2), padding='same'),
    l.Flatten(),
    sparsity.prune_low_magnitude(l.Dense(1024, activation='relu'),
                                 **pruning_params),
    l.Dropout(0.4),
    sparsity.prune_low_magnitude(l.Dense(num_classes, activation='softmax'),
                                 **pruning_params)
])

... I had to restart Jupyter kernel to get rid of further errors, like Conv2D has no attribute 'kernel', as seen at GitHub:

tf.enable_eager_execution must be called at program startup. #18304