Fix docstring

facebookresearch · Dec 7, 2024 · 1bbaa2f · 1bbaa2f
1 parent 7e20cf3
commit 1bbaa2f
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/distributed_shampoo/distributed_shampoo.py b/distributed_shampoo/distributed_shampoo.py
@@ -272,7 +272,7 @@ class DistributedShampoo(torch.optim.Optimizer):
         preconditioner_config (PreconditionerConfig): Configuration for preconditioner computation.
             If this field is an instance ShampooPreconditionerConfig, Shampoo uses the root inverse of the preconditioner.
             If this field is an instance EigenvalueCorrectedShampooPreconditionerConfig Shampoo uses corrected the eigenvalues/running Adam in the eigenbasis of preconditioner.
-            (Default: DefaultEigenConfig)
+            (Default: DefaultShampooConfig)
 
     """