Merge pull request #575 from benkli01:toupstream/revert_tidy_up

PiperOrigin-RevId: 338340110

Author: tensorflower-gardener

tensorflow_model_optimization/python/core/clustering/keras/cluster_wrapper.py

@@ -326,11 +326,11 @@ def trainable(self, value):
 
   @property
   def trainable_weights(self):
-    return self.layer.trainable_weights
+    return self.layer.trainable_weights + self._trainable_weights
 
   @property
   def non_trainable_weights(self):
-    return self.layer.non_trainable_weights
+    return self.layer.non_trainable_weights + self._non_trainable_weights
 
   @property
   def updates(self):