FIX: weightnorm variables (#219)

* FIX: weightnorm variables

Author: seanpmorgan

tensorflow_addons/layers/wrappers.py

@@ -88,6 +88,7 @@ def build(self, input_shape):
 
         super(WeightNormalization, self).build()
 
+    @tf.function
     def call(self, inputs):
         """Call `Layer`"""
         if not self._initialized:
@@ -143,9 +144,9 @@ def _data_dep_init(self, inputs):
             scale_init = 1. / tf.math.sqrt(v_init + 1e-10)
 
         # Assign data dependent init values
-        self.g.assign(self.g * scale_init)
+        self.g = self.g * scale_init
         if hasattr(self.layer, 'bias'):
-            self.layer.bias.assign(-m_init * scale_init)
+            self.layer.bias = -m_init * scale_init
         self.layer.activation = existing_activation
 
     def get_config(self):

tensorflow_addons/layers/wrappers_test.py

@@ -75,7 +75,7 @@ def test_weightnorm_conv2d(self):
 
         self.assertTrue(hasattr(model.layers[0], 'g'))
 
-    def test_weightnorm_tflayers(self):
+    def test_weightnorm_applylayer(self):
         images = tf.random.uniform((2, 4, 4, 3))
         wn_wrapper = wrappers.WeightNormalization(
             tf.keras.layers.Conv2D(32, [2, 2]), input_shape=(4, 4, 3))