Fix unit test failure caused by the keras initializer change.

PiperOrigin-RevId: 453982185

official/nlp/modeling/layers/tn_expand_condense_test.py

@@ -66,7 +66,7 @@ class TNLayerTest(tf.test.TestCase, parameterized.TestCase):
   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 @@ class TNLayerTest(tf.test.TestCase, parameterized.TestCase):
 
   @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(