Internal

PiperOrigin-RevId: 483746508
google · Oct 26, 2022 · c07ca02 · c07ca02
1 parent 8768fb9
commit c07ca02
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/baselines/jft/test_utils.py b/baselines/jft/test_utils.py
@@ -114,7 +114,7 @@ def get_config(
 
   config.lr = ml_collections.ConfigDict()
   config.lr.base = 0.1
-  config.lr.warmup_steps = 1
+  config.lr.warmup_steps = 1 if total_steps > 1 else 0
   config.lr.decay_type = 'linear'
   config.lr.linear_end = 1e-3