fix: add max_memory to _init_infer_auto_device_map's return statement (…

…#3279)
huggingface · Dec 13, 2024 · 200c9eb · 200c9eb
1 parent 7b2edc0
commit 200c9eb
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/src/accelerate/utils/modeling.py b/src/accelerate/utils/modeling.py
@@ -1101,6 +1101,7 @@ def _init_infer_auto_device_map(
     special_dtypes: Optional[Dict[str, Union[str, torch.device]]] = None,
 ) -> Tuple[
     List[Union[int, str]],
+    Dict[Union[int, str], Union[int, str]],
     List[Union[int, str]],
     List[int],
     Dict[str, int],
@@ -1147,6 +1148,7 @@ def _init_infer_auto_device_map(
 
     return (
         devices,
+        max_memory,
         main_devices,
         gpus,
         module_sizes,
@@ -1356,6 +1358,7 @@ def infer_auto_device_map(
     # Initialize the variables
     (
         devices,
+        max_memory,
         main_devices,
         gpus,
         module_sizes,