sparsity.prune_low_magnitude fails with mixed precision policy mixed_float16

Describe the bug
When using tf.keras.mixed_precision.experimental.Policy("mixed_float16", loss_scale="dynamic") the sparsity.prune_low_magnitude fails in tensor type conversion with the error Tensor conversion requested dtype float32 for Tensor with dtype float16: <tf.Tensor 'pruning_ops/Cast_2:0' shape=(3, 3, 1, 12) dtype=float16>. Things work perfectly fine when precision is set to the default float32. Looks like some piece of code is not properly respecting the dtype.

System information

TensorFlow installed from (source or binary): pip3

TensorFlow version: 2.2.0

TensorFlow Model Optimization version:

Python version: 3.6

Describe the expected behavior

The prune_low_magnitude should work with layers using mixed_float16 policy.

Describe the current behavior

Throws error described above.

Code to reproduce the issue
See this colab link.

Screenshots
If applicable, add screenshots to help explain your problem.

Additional context
Add any other context about the problem here.