From 200c9eb7833cfa505907f6f224ebf5a275aa6d92 Mon Sep 17 00:00:00 2001
From: Chen <xiaokunchen@outlook.com>
Date: Fri, 13 Dec 2024 10:47:33 -0500
Subject: [PATCH] fix: add max_memory to _init_infer_auto_device_map's return
 statement (#3279)

---
 src/accelerate/utils/modeling.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/src/accelerate/utils/modeling.py b/src/accelerate/utils/modeling.py
index 73ada66544f..2b7e5fcaa19 100644
--- 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,