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

PiperOrigin-RevId: 387447417

official/nlp/modeling/layers/spectral_normalization.py

@@ -106,16 +106,19 @@ class SpectralNormalization(tf.keras.layers.Wrapper):
 
   def call(self, inputs, *, training=None):
     training = self.do_power_iteration if training is None else training
-    u_update_op, v_update_op, w_update_op = self.update_weights(
-        training=training)
-    output = self.layer(inputs)
-    w_restore_op = self.restore_weights()
-
-    # Register update ops.
-    self.add_update(u_update_op)
-    self.add_update(v_update_op)
-    self.add_update(w_update_op)
-    self.add_update(w_restore_op)
+    if training:
+      u_update_op, v_update_op, w_update_op = self.update_weights(
+          training=training)
+      output = self.layer(inputs)
+      w_restore_op = self.restore_weights()
+
+      # Register update ops.
+      self.add_update(u_update_op)
+      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