From f20130869629814d7ff6fc45b503b46c26411f86 Mon Sep 17 00:00:00 2001 From: David B Date: Sun, 10 Mar 2024 19:42:26 +0000 Subject: [PATCH] Fix pydantic configuration for the max_memory input --- src/axolotl/utils/config/models/input/v0_4_1/__init__.py | 2 +- 1 file changed, 1 insertion(+), 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 index 79fffe9cd6..b0888c3b1d 100644 --- 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