Unsets profile_duration_ms in Flax ResNet example to prioritize `nu…

…m_profile_steps`. PiperOrigin-RevId: 648828498
google · Jul 2, 2024 · 3898afd · 3898afd
1 parent 37123d5
commit 3898afd
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/examples/imagenet/train.py b/examples/imagenet/train.py
@@ -372,7 +372,8 @@ 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, logdir=workdir, profile_duration_ms=None)]
   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):