pruned pretrained mobilenet_v1_1.0_224, mode.fit() error
@PPParticle Thanks for reporting the issue.

Sorry for the delayed response.

As per the documentation, [prune_low_magnitude](https://www.tensorflow.org/model_optimization/api_docs/python/tfmot/sparsity/keras/prune_low_magnitude) expects pruning schedule, but a dictionary of pruning schedule is given which can be resolved by either directly passing pruing schedule or getting pruning schedule from dictionary. Also, the callbacks are not  called during pruned `model.fit`. 

Please find the working [gist](https://colab.research.google.com/gist/pjpratik/3b25bca188e4b09526addf1418c197b0/59679.ipynb) and let us know if it helps.

The weights are not loaded as they are not provided.

Thanks! > 
Thanks so much for your reply which has solved my problems !! However, I haven't figure out the reason why there are still so many non-trainable parameters after `model.fit()` with `model_for_pruning.summary()`. Maybe I should look up TF guidebook! All in all, thank you so much again !! @PPParticle Please feel free to close the issue if it is resolved. Thanks! Are you satisfied with the resolution of your issue?
<a href="https://docs.google.com/forms/d/e/1FAIpQLSfaP12TRhd9xSxjXZjcZFNXPGk4kc1-qMdv3gc6bEP90vY1ew/viewform?entry.85265664=Yes&entry.2137816233=https://github.com/tensorflow/tensorflow/issues/59679">Yes</a>
<a href="https://docs.google.com/forms/d/e/1FAIpQLSfaP12TRhd9xSxjXZjcZFNXPGk4kc1-qMdv3gc6bEP90vY1ew/viewform?entry.85265664=No&entry.2137816233=https://github.com/tensorflow/tensorflow/issues/59679">No</a>

### Issue Type

Bug

### Have you reproduced the bug with TF nightly?

No

### Source

binary

### Tensorflow Version

v2.11.0

### Custom Code

No

### OS Platform and Distribution

WSL2 Linux Ubuntu 20.04

### Mobile device

_No response_

### Python version

python 3.8.10

### Bazel version

_No response_

### GCC/Compiler version

_No response_

### CUDA/cuDNN version

_No response_

### GPU model and memory

_No response_

### Current Behaviour?

```shell
A bug happened!
I wanted to prun pretrained mobilenet_v1 with model_weight mobilenet_1_0_224. After I pruned the pretrained model and start to retrain pruned_model, something was wrong.
```


### Standalone code to reproduce the issue

```shell
This is code

from matplotlib import pyplot as plt
import tensorflow_model_optimization as tfmot
import numpy as np
import os
import tensorflow as tf
from tensorflow import keras
import tensorflow_datasets as tfds
import tensorflow_model_optimization as tfmot
import tempfile

_URL = 'https://storage.googleapis.com/mledu-datasets/cats_and_dogs_filtered.zip'
path_to_zip = tf.keras.utils.get_file('cats_and_dogs.zip', origin=_URL, extract=True)
PATH = os.path.join(os.path.dirname(path_to_zip), 'cats_and_dogs_filtered')

train_path = os.path.join(PATH, 'train')
validation_path = os.path.join(PATH, 'validation')

IMG_SIZE = (224, 224)

train_dataset = tf.keras.utils.image_dataset_from_directory(train_path, shuffle=True, image_size=IMG_SIZE)
validation_dataset = tf.keras.utils.image_dataset_from_directory(validation_path, shuffle=True, image_size=IMG_SIZE)
 
val_batches = tf.data.experimental.cardinality(validation_dataset) # check data number of one batch
test_dataset = validation_dataset.take(val_batches // 5) 
val_dataset = validation_dataset.skip(val_batches // 5)

num_images = tf.data.experimental.cardinality(test_dataset)
end_step = np.ceil(num_images / 32).astype(np.int32) * 2

# load base model and freeze the base model
model = tf.keras.applications.MobileNet()
model.load_weights(r'/home/user/.keras/models/mobilenet_1_0_224_tf.h5')

model.compile(optimizer='adam',
              loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
              metrics=['accuracy'])

model.fit(test_dataset, epochs=2, validation_data=val_dataset)

_, keras_file = tempfile.mkstemp('.h5')
tf.keras.models.save_model(model, keras_file, include_optimizer=False)
print('Saved baseline model to:', keras_file)
model.save("mobilenet_no_pruning.h5")
# loss, accuracy = model.evaluate(validation_dataset)              
# # model_for_pruning.summary() # trainable parameters + non-trainable mask parameters
# print("initial loss: {:.2f}".format(loss))
# print("initial accuracy: {:.2f}".format(accuracy))

# model applied pruning  
pruning_schedule = {
      'pruning_schedule': tfmot.sparsity.keras.PolynomialDecay(0.5,0.9,begin_step=0, end_step=end_step)
}        
model_for_pruning = tfmot.sparsity.keras.prune_low_magnitude(model, pruning_schedule=pruning_schedule)

# create a new model
model_for_pruning.compile(optimizer='adam',
              loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
              metrics=['accuracy'])

logdir = tempfile.mkdtemp()
callbacks = [
  tfmot.sparsity.keras.UpdatePruningStep(), # pruning parameters
  tfmot.sparsity.keras.PruningSummaries(log_dir=logdir),
]
print("begin fit")
# model_for_pruning.fit(test_dataset, batch_size=32, epochs=2, validation_data=val_dataset, callbacks=callbacks)
model_for_pruning.fit(test_dataset, batch_size=32, epochs=2, validation_data=val_dataset)
loss, accuracy = model_for_pruning.evaluate(validation_dataset)              
# # model_for_pruning.summary() # trainable parameters + non-trainable mask parameters
print("initial loss: {:.2f}".format(loss))
print("initial accuracy: {:.2f}".format(accuracy))
```
```


### Relevant log output

The pretrained mode traing is as usual. The errors happens retrained pruned model. The errors are 

```shell
Found 2000 files belonging to 2 classes.
Found 1000 files belonging to 2 classes.
Epoch 1/2
/usr/local/lib/python3.8/dist-packages/keras/backend.py:5585: UserWarning: "`sparse_categorical_crossentropy` received `from_logits=True`, but the `output` argument was produced by a Softmax activation and thus does not represent logits. Was this intended?
  output, from_logits = _get_logits(
6/6 [==============================] - 26s 4s/step - loss: 4.2002 - accuracy: 0.5365 - val_loss: 28.3055 - val_accuracy: 0.0000e+00
Epoch 2/2
6/6 [==============================] - 19s 3s/step - loss: 0.3738 - accuracy: 0.9479 - val_loss: 21.0645 - val_accuracy: 0.0223
Saved baseline model to: /tmp/tmp68woeyxc.h5
WARNING:tensorflow:From /home/qiao_sh_pudong/.local/lib/python3.8/site-packages/tensorflow/python/autograph/pyct/static_analysis/liveness.py:83: Analyzer.lamba_check (from tensorflow.python.autograph.pyct.static_analysis.liveness) is deprecated and will be removed after 2023-09-23.
Instructions for updating:
Lambda fuctions will be no more assumed to be used in the statement where they are used, or at least in the same block. https://github.com/tensorflow/tensorflow/issues/56089
begin fit
Epoch 1/2
Traceback (most recent call last):
  File "fine-tune_mobilenet_v1_with_pruning.py", line 85, in <module>
    model_for_pruning.fit(test_dataset, batch_size=32, epochs=2, validation_data=val_dataset)
  File "/usr/local/lib/python3.8/dist-packages/keras/utils/traceback_utils.py", line 70, in error_handler
    raise e.with_traceback(filtered_tb) from None
  File "/tmp/__autograph_generated_filefab8qmeh.py", line 15, in tf__train_function
    retval_ = ag__.converted_call(ag__.ld(step_function), (ag__.ld(self), ag__.ld(iterator)), None, fscope)
  File "/tmp/__autograph_generated_filerwugbiee.py", line 71, in tf__call
    update_mask = ag__.converted_call(ag__.ld(utils).smart_cond, (ag__.ld(training), ag__.ld(add_update), ag__.ld(no_op)), None, fscope)
  File "/tmp/__autograph_generated_filepmfh1utc.py", line 37, in tf__smart_cond
    ag__.if_stmt(ag__.converted_call(ag__.ld(isinstance), (ag__.ld(pred), ag__.ld(variables).Variable), None, fscope), if_body, else_body, get_state, set_state, ('do_return', 'retval_'), 2)
  File "/tmp/__autograph_generated_filepmfh1utc.py", line 33, in else_body
    retval_ = ag__.converted_call(ag__.ld(smart_module).smart_cond, (ag__.ld(pred),), dict(true_fn=ag__.ld(true_fn), false_fn=ag__.ld(false_fn), name=ag__.ld(name)), fscope)
  File "/tmp/__autograph_generated_filerwugbiee.py", line 48, in add_update
    with ag__.ld(tf).control_dependencies([ag__.ld(self).pruning_obj.conditional_mask_update()]):
  File "/home/qiao_sh_pudong/.local/lib/python3.8/site-packages/tensorflow_model_optimization/python/core/sparsity/keras/pruning_impl.py", line 310, in conditional_mask_update
    return tf.distribute.get_replica_context().merge_call(
  File "/tmp/__autograph_generated_filepd0d14ah.py", line 63, in tf__mask_update_distributed
    retval_ = ag__.converted_call(ag__.ld(tf).cond, (ag__.converted_call(ag__.ld(maybe_update_masks), (), None, fscope), ag__.ld(update_distributed), ag__.ld(no_update)), None, fscope)
  File "/tmp/__autograph_generated_file1681ktv2.py", line 37, in tf__maybe_update_masks
    ag__.if_stmt(ag__.ld(self)._sparsity_m_by_n, if_body, else_body, get_state, set_state, ('do_return', 'retval_'), 2)
  File "/tmp/__autograph_generated_file1681ktv2.py", line 33, in else_body
    retval_ = ag__.converted_call(ag__.ld(self)._pruning_schedule, (ag__.converted_call(ag__.ld(self)._step_fn, (), None, fscope),), None, fscope)[0]
TypeError: in user code:

    File "/usr/local/lib/python3.8/dist-packages/keras/engine/training.py", line 1249, in train_function  *
        return step_function(self, iterator)
    File "/usr/local/lib/python3.8/dist-packages/keras/engine/training.py", line 1233, in step_function  **
        outputs = model.distribute_strategy.run(run_step, args=(data,))
    File "/usr/local/lib/python3.8/dist-packages/keras/engine/training.py", line 1222, in run_step  **
        outputs = model.train_step(data)
    File "/usr/local/lib/python3.8/dist-packages/keras/engine/training.py", line 1023, in train_step
        y_pred = self(x, training=True)
    File "/usr/local/lib/python3.8/dist-packages/keras/utils/traceback_utils.py", line 70, in error_handler
        raise e.with_traceback(filtered_tb) from None
    File "/tmp/__autograph_generated_filerwugbiee.py", line 71, in tf__call
        update_mask = ag__.converted_call(ag__.ld(utils).smart_cond, (ag__.ld(training), ag__.ld(add_update), ag__.ld(no_op)), None, fscope)
    File "/tmp/__autograph_generated_filepmfh1utc.py", line 37, in tf__smart_cond
        ag__.if_stmt(ag__.converted_call(ag__.ld(isinstance), (ag__.ld(pred), ag__.ld(variables).Variable), None, fscope), if_body, else_body, get_state, set_state, ('do_return', 'retval_'), 2)
    File "/tmp/__autograph_generated_filepmfh1utc.py", line 33, in else_body
        retval_ = ag__.converted_call(ag__.ld(smart_module).smart_cond, (ag__.ld(pred),), dict(true_fn=ag__.ld(true_fn), false_fn=ag__.ld(false_fn), name=ag__.ld(name)), fscope)
    File "/tmp/__autograph_generated_filerwugbiee.py", line 48, in add_update
        with ag__.ld(tf).control_dependencies([ag__.ld(self).pruning_obj.conditional_mask_update()]):
    File "/home/qiao_sh_pudong/.local/lib/python3.8/site-packages/tensorflow_model_optimization/python/core/sparsity/keras/pruning_impl.py", line 310, in conditional_mask_update
        return tf.distribute.get_replica_context().merge_call(
    File "/tmp/__autograph_generated_filepd0d14ah.py", line 63, in tf__mask_update_distributed
        retval_ = ag__.converted_call(ag__.ld(tf).cond, (ag__.converted_call(ag__.ld(maybe_update_masks), (), None, fscope), ag__.ld(update_distributed), ag__.ld(no_update)), None, fscope)
    File "/tmp/__autograph_generated_file1681ktv2.py", line 37, in tf__maybe_update_masks
        ag__.if_stmt(ag__.ld(self)._sparsity_m_by_n, if_body, else_body, get_state, set_state, ('do_return', 'retval_'), 2)
    File "/tmp/__autograph_generated_file1681ktv2.py", line 33, in else_body
        retval_ = ag__.converted_call(ag__.ld(self)._pruning_schedule, (ag__.converted_call(ag__.ld(self)._step_fn, (), None, fscope),), None, fscope)[0]

    TypeError: Exception encountered when calling layer 'prune_low_magnitude_conv1' (type PruneLowMagnitude).
    
    in user code:
    
        File "/home/qiao_sh_pudong/.local/lib/python3.8/site-packages/tensorflow_model_optimization/python/core/sparsity/keras/pruning_wrapper.py", line 268, in add_update  *
            with tf.control_dependencies(
        File "/home/qiao_sh_pudong/.local/lib/python3.8/site-packages/tensorflow_model_optimization/python/core/keras/utils.py", line 54, in smart_cond  *
            pred, true_fn=true_fn, false_fn=false_fn, name=name)
        File "/home/qiao_sh_pudong/.local/lib/python3.8/site-packages/tensorflow_model_optimization/python/core/sparsity/keras/pruning_impl.py", line 307, in mask_update_distributed  *
            return tf.cond(maybe_update_masks(), update_distributed, no_update)
        File "/home/qiao_sh_pudong/.local/lib/python3.8/site-packages/tensorflow_model_optimization/python/core/sparsity/keras/pruning_impl.py", line 264, in maybe_update_masks  *
            return self._pruning_schedule(self._step_fn())[0]
    
        TypeError: '_DictWrapper' object is not callable
    
    
    Call arguments received by layer 'prune_low_magnitude_conv1' (type PruneLowMagnitude):
      • inputs=tf.Tensor(shape=(None, 224, 224, 3), dtype=float32)
      • training=True
      • kwargs=<class 'inspect._empty'>
```
