log_step_hint

mlcommons · Sep 28, 2023 · c3cf664 · c3cf664
1 parent 1ea08b0
commit c3cf664
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/reference_algorithms/target_setting_algorithms/jax_nadamw.py b/reference_algorithms/target_setting_algorithms/jax_nadamw.py
@@ -7,6 +7,7 @@
 import jax
 import jax.numpy as jnp
 import optax
+from absl import logging
 
 from algorithmic_efficiency import spec
 from reference_algorithms.target_setting_algorithms import cosine_warmup
@@ -152,6 +153,7 @@ def init_optimizer_state(workload: spec.Workload,
   del rng
 
   target_setting_step_hint = int(0.75 * workload.step_hint)
+  logging.info(f'target setting step hint: {target_setting_step_hint}')
   lr_schedule_fn = cosine_warmup.jax_cosine_warmup(target_setting_step_hint,
                                                    hyperparameters)
 
@@ -168,4 +170,4 @@ def init_optimizer_state(workload: spec.Workload,
       weight_decay=hyperparameters.weight_decay)
   optimizer_state = opt_init_fn(params_zeros_like)
 
-  return jax_utils.replicate(optimizer_state), opt_update_fn, lr_schedule_fn
+  return jax_utils.replicate(optimizer_state), opt_update_fn, opt_update_fn