Repository: huggingface/accelerate
Issue category: Bug Report

Issue:
Calling infer_auto_device_map() with max_memory=None throws an error in version 1.2.0
### System Info

```Shell
accelerate==1.2.0
```

### Reproduction

Bug is from this commit:
https://github.com/huggingface/accelerate/commit/d7b1b368e9f484a18636a71600566b757d5cf87e

`max_memory` initialization was moved into `_init_infer_auto_device_map`, which does not return the `max_memory` value.

So if max_memory=None is passed to `infer_auto_device_map` (the default value), then it will still be None at line 1415:
https://github.com/huggingface/accelerate/blob/cb8b7c637a8588668c52bd306f9b2828f69d9585/src/accelerate/utils/modeling.py#L1415

Leading to error: TypeError: 'NoneType' object is not subscriptable

### Expected behavior

max_memory=None when passed to `infer_auto_device_map` does not throw an error.


Hints:
@Nech-C 
Sorry for the oversight. I will fix it ASAP. Thanks for pointing it out!

Patch:
diff --git a/src/accelerate/utils/modeling.py b/src/accelerate/utils/modeling.py
index 5f88e54e3c9..806f930acaa 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,


Test patch: