Fix pydantic configuration for the max_memory input (#1385) [skip ci]

* Fix pydantic configuration for the max_memory input

* chore: lint

---------

Co-authored-by: Wing Lian <[email protected]>

src/axolotl/utils/config/models/input/v0_4_1/__init__.py

@@ -512,7 +512,9 @@ class AxolotlInputConfig(
 
     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