turn on preconditioning

mlcommons · Feb 10, 2024 · 73147fd · 73147fd
1 parent 1307931
commit 73147fd
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/reference_algorithms/paper_baselines/shampoo/jax/distributed_shampoo.py b/reference_algorithms/paper_baselines/shampoo/jax/distributed_shampoo.py
@@ -1318,7 +1318,7 @@ def distributed_shampoo(
     ###
     inverse_failure_threshold=0.1,
     moving_average_for_momentum=True,
-    skip_preconditioning_dim_size_gt=0,
+    skip_preconditioning_dim_size_gt=4096,
     clip_by_scaled_gradient_norm=None,
     precision=lax.Precision.HIGHEST,
     tensordot_precision=None,