pruned_model = setup_pruned_model() loss = tf.keras.losses.categorical_crossentropy optimizer = keras.optimizers.Adam() log_dir = tempfile.mkdtemp() # This is all not boilerplate. pruned_model.optimizer = optimizer step_callback = tfmot.sparsity.keras.UpdatePruningStep() step_callback.set_model(pruned_model) log_callback = tfmot.sparsity.keras.PruningSummaries(log_dir=log_dir) # optional Tensorboard logging. log_callback.set_model(pruned_model) step_callback.on_train_begin() for _ in range(3): # only one batch given batch_size = 20 and input shape. step_callback.on_train_batch_begin(batch=unused_arg) inp = np.reshape(x_train, [self._BATCH_SIZE, 10]) # original shape: from [10]. with tf.GradientTape() as tape: logits = pruned_model(inp, training=True) loss_value = loss(y_train, logits) grads = tape.gradient(loss_value, pruned_model.trainable_variables) optimizer.apply_gradients(zip(grads, pruned_model.trainable_variables)) step_callback.on_epoch_end(batch=unused_arg) log_callback.on_epoch_end(batch=unused_arg) ...