Fix unit test failure caused by the keras initializer change.

PiperOrigin-RevId: 453982185
tly18 · Jun 9, 2022 · 44848e0 · 44848e0
1 parent a506d3a
commit 44848e0
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/official/nlp/modeling/layers/tn_expand_condense_test.py b/official/nlp/modeling/layers/tn_expand_condense_test.py
@@ -66,7 +66,7 @@ def test_keras_layer(self, input_dim, proj_multiple):
   def test_train(self, input_dim, proj_multiple):
     data = np.random.randint(10, size=(100, input_dim))
     model = self._build_model(data, proj_multiple)
-    tf.random.set_seed(0)
+    tf.keras.utils.set_random_seed(0)
 
     model.compile(
         optimizer='adam', loss='binary_crossentropy', metrics=['accuracy'])
@@ -81,7 +81,7 @@ def test_train(self, input_dim, proj_multiple):
 
   @parameterized.parameters((768, 6), (1024, 2))
   def test_weights_change(self, input_dim, proj_multiple):
-    tf.random.set_seed(0)
+    tf.keras.utils.set_random_seed(0)
     data = np.random.randint(10, size=(100, input_dim))
     model = self._build_model(data, proj_multiple)
     model.compile(