Merge branch 'configs-refactor' into fail-counter

facebookresearch · Dec 10, 2024 · d53ffc6 · d53ffc6
2 parents 709ab1c + 273e0a1
commit d53ffc6
Showing 1 changed file with 7 additions and 3 deletions.
diff --git a/distributed_shampoo/shampoo_types.py b/distributed_shampoo/shampoo_types.py
@@ -8,7 +8,7 @@
 """
 
 import enum
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 
 import torch
 
@@ -109,7 +109,9 @@ class ShampooPreconditionerConfig(PreconditionerConfig):
 
     """
 
-    amortized_computation_config: RootInvConfig = DefaultEigenConfig
+    amortized_computation_config: RootInvConfig = field(
+        default_factory=lambda: DefaultEigenConfig
+    )
 
 
 DefaultShampooConfig = ShampooPreconditionerConfig()
@@ -124,7 +126,9 @@ class EigenvalueCorrectedShampooPreconditionerConfig(PreconditionerConfig):
 
     """
 
-    amortized_computation_config: EigenvectorConfig = DefaultEighConfig
+    amortized_computation_config: EigenvectorConfig = field(
+        default_factory=lambda: DefaultEighConfig
+    )
 
 
 DefaultEigenvalueCorrectedShampooConfig = (