Fix pydantic configuration for the max_memory input

axolotl-ai-cloud · Mar 10, 2024 · f201308 · f201308
1 parent 9b6ee83
commit f201308
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/axolotl/utils/config/models/input/v0_4_1/__init__.py b/src/axolotl/utils/config/models/input/v0_4_1/__init__.py
@@ -511,7 +511,7 @@ class Config:
 
     neftune_noise_alpha: Optional[float] = None
 
-    max_memory: Optional[Union[int, str]] = None
+    max_memory: Optional[Dict[Union[int, Literal['cpu','disk']],Union[int, str]]] = None
     gpu_memory_limit: Optional[Union[int, str]] = None
 
     chat_template: Optional[Union[Literal["chatml", "inst"], ChatTemplate]] = None