SpectralNormalization: Avoid errors in export by only applying update ops when training=True.

PiperOrigin-RevId: 387447417

official/nlp/modeling/layers/spectral_normalization.py

@@ -106,6 +106,7 @@ class SpectralNormalization(tf.keras.layers.Wrapper):
 
   def call(self, inputs, *, training=None):
     training = self.do_power_iteration if training is None else training
+    if training:
       u_update_op, v_update_op, w_update_op = self.update_weights(
           training=training)
       output = self.layer(inputs)
@@ -116,6 +117,8 @@ class SpectralNormalization(tf.keras.layers.Wrapper):
       self.add_update(v_update_op)
       self.add_update(w_update_op)
       self.add_update(w_restore_op)
+    else:
+      output = self.layer(inputs)
 
     return output