Propogate training arg to inner layer in QuantizeWrapper

PiperOrigin-RevId: 321031035

Author: nutsiepully

tensorflow_model_optimization/python/core/quantization/keras/quantize_wrapper.py

@@ -26,6 +26,7 @@
 from __future__ import division
 from __future__ import print_function
 
+import inspect
 import tensorflow as tf
 
 # TODO(b/139939526): move to public API.
@@ -159,7 +160,11 @@ def call(self, inputs, training=None):
     self.quantize_config.set_quantize_activations(self.layer,
                                                   self._quantize_activations)
 
-    outputs = self.layer.call(inputs)
+    args = inspect.getfullargspec(self.layer.call).args
+    if 'training' in args:
+      outputs = self.layer.call(inputs, training=training)
+    else:
+      outputs = self.layer.call(inputs)
 
     if not self._output_quantizers:
       return outputs