'prune_low_magnitude' of 'keras.layers.convolutional.Conv2D' returns None #163
Description
When trying to prune the Cov2D layer in a custom network the below error is observed.
TypeError: 'NoneType' object is not callable
Below is a snippet of the code:
from tensorflow_model_optimization.sparsity import keras as sparsity
from keras.layers.convolutional import Conv2D, Conv2DTranspose
pruning_prm = {'pruning_schedule': sparsity.PolynomialDecay(initial_sparsity=0, final_sparsity=0.5,begin_step=15000, end_step=75000, frequency=100)}
x = Conv2D(filters=n_filters, kernel_size=(kernel_size, kernel_size), kernel_initializer="he_normal",
padding="same")
x = sparsity.prune_low_magnitude(x, **pruning_prm) --> returns None
Information on versions of different modules:
Keras 2.3.0
tensorflow-estimator (2.0.0)
tensorflow-gpu (2.0.0)
tensorflow-hub (0.6.0)
tensorflow-model-optimization (0.1.3)
tensorflow-probability (0.8.0rc0)
tensorflow-serving-api-gpu (1.14.0)
Any pointers/leads is much appreciated.