In pruning_wrapper, propagate training boolean to the wrapped layer.

PiperOrigin-RevId: 303829300

Author: Suyog Gupta

tensorflow_model_optimization/python/core/sparsity/keras/pruning_wrapper.py

@@ -19,6 +19,7 @@
 from __future__ import division
 from __future__ import print_function
 
+import inspect
 # import g3
 import numpy as np
 import tensorflow as tf
@@ -255,6 +256,12 @@ def no_op():
     # self.add_update does nothing during eager execution.
     self.add_update(self.pruning_obj.weight_mask_op())
 
+    args = inspect.getargspec(self.layer.call)[0]
+    # Propagate the training bool to the underlying layer if it accepts
+    # training as an arg.
+    if 'training' in args:
+      return self.layer.call(inputs, training=training)
+
     return self.layer.call(inputs)
 
   def compute_output_shape(self, input_shape):