Adds save_steps to fix squad

PiperOrigin-RevId: 275994710

official/nlp/xlnet/common_flags.py

@@ -142,5 +142,5 @@ flags.DEFINE_integer(
     default=16,
     help="Size of the test batch across all hosts.")
 flags.DEFINE_integer(
-    "save_steps", default=None, help="Number of steps for saving checkpoint.")
+    "save_steps", default=1000, help="Number of steps for saving checkpoint.")
 FLAGS = flags.FLAGS

official/nlp/xlnet/run_classifier.py

@@ -184,7 +184,7 @@ def main(unused_argv):
       optimizer=optimizer,
       learning_rate_fn=learning_rate_fn,
       model_dir=FLAGS.model_dir,
-      save_steps=1000)
+      save_steps=FLAGS.save_steps)
 
 
 if __name__ == "__main__":

official/nlp/xlnet/run_squad.py

@@ -288,7 +288,8 @@ def main(unused_argv):
       steps_per_loop=steps_per_loop,
       optimizer=optimizer,
       learning_rate_fn=learning_rate_fn,
-      model_dir=FLAGS.model_dir)
+      model_dir=FLAGS.model_dir,
+      save_steps=FLAGS.save_steps)
 
 
 if __name__ == "__main__":