Add test for saving

tensorflow · Nov 12, 2019 · 833f807 · 833f807
1 parent d29c51d
commit 833f807
Showing 1 changed file with 11 additions and 0 deletions.
diff --git a/tensorflow_addons/layers/wrappers_test.py b/tensorflow_addons/layers/wrappers_test.py
@@ -81,6 +81,17 @@ def test_weightnorm_with_time_dist(self):
         out = tf.keras.layers.TimeDistributed(b)(inputs)
         model = tf.keras.Model(inputs, out)
 
+    def test_save_file_h5(self):
+        conv = tf.keras.layers.Conv1D(1, 1)
+        wn_conv = wrappers.WeightNormalization(conv)
+        model = tf.keras.Sequential(layers=[wn_conv])
+        model.build([1, 2, 3])
+        model.save_weights('/tmp/model.h5')
+
+        import os
+        os.remove('/tmp/model.h5')
+        # TODO: Find a better way to test this
+
 
 if __name__ == "__main__":
     tf.test.main()