Fix defaults with default_factory

facebookresearch · Dec 10, 2024 · 273e0a1 · 273e0a1
1 parent d3a10af
commit 273e0a1
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
 
@@ -101,7 +101,9 @@ class ShampooPreconditionerConfig(PreconditionerConfig):
 
     """
 
-    amortized_computation_config: RootInvConfig = DefaultEigenConfig
+    amortized_computation_config: RootInvConfig = field(
+        default_factory=lambda: DefaultEigenConfig
+    )
 
 
 DefaultShampooConfig = ShampooPreconditionerConfig()
@@ -116,7 +118,9 @@ class EigenvalueCorrectedShampooPreconditionerConfig(PreconditionerConfig):
 
     """
 
-    amortized_computation_config: EigenvectorConfig = DefaultEighConfig
+    amortized_computation_config: EigenvectorConfig = field(
+        default_factory=lambda: DefaultEighConfig
+    )
 
 
 DefaultEigenvalueCorrectedShampooConfig = (