set profile_duration_ms = None as in periodic_actions there's default…

… value for both num_profile_steps and profile_duration_ms, and the profile stopping condition is when both num_profile_steps and profile_duration_ms are satisfied, so setting profile_duration_ms=None so that the passed num_profile_steps value gets used

PiperOrigin-RevId: 657605586

examples/imagenet/train.py

@@ -372,7 +372,11 @@ def train_and_evaluate(
   train_metrics = []
   hooks = []
   if jax.process_index() == 0 and config.profile:
-    hooks += [periodic_actions.Profile(num_profile_steps=3, logdir=workdir)]
+    hooks += [
+        periodic_actions.Profile(
+            num_profile_steps=3, profile_duration_ms=None, logdir=workdir
+        )
+    ]
   train_metrics_last_t = time.time()
   logging.info('Initial compilation, this might take some minutes...')
   for step, batch in zip(range(step_offset, num_steps), train_iter):