diff --git a/src/accelerate/utils/modeling.py b/src/accelerate/utils/modeling.py
index 1760e72bc..c14833557 100644
--- a/src/accelerate/utils/modeling.py
+++ b/src/accelerate/utils/modeling.py
@@ -408,8 +408,9 @@ def get_balanced_memory(
     per_gpu += buffer
 
     max_memory = get_max_memory(max_memory)
+    last_gpu = max(i for i in max_memory if isinstance(i, int) and max_memory[i] > 0)
     # The last device is left with max_memory just in case the buffer is not enough.
-    for i in range(len(max_memory) - 1):
+    for i in range(last_gpu):
         max_memory[i] = min(0 if low_zero and i == 0 else per_gpu, max_memory[i])
 
     if low_zero:
diff --git a/tests/test_modeling_utils.py b/tests/test_modeling_utils.py
index 02ab867d7..16243d5e1 100644
--- a/tests/test_modeling_utils.py
+++ b/tests/test_modeling_utils.py
@@ -376,6 +376,10 @@ def test_get_balanced_memory(self):
         max_memory = get_balanced_memory(model, max_memory={0: 300, 1: 500})
         self.assertDictEqual({0: 215, 1: 500}, max_memory)
 
+        # Last device always get max memory to give more buffer, even if CPU is provided
+        max_memory = get_balanced_memory(model, max_memory={0: 300, "cpu": 1000})
+        self.assertDictEqual({0: 300, "cpu": 1000}, max_memory)
+
         # If we set a device to 0, it's not counted.
         max_memory = get_balanced_memory(model, max_memory={0: 0, 1: 300, 2: 300})
         self.assertDictEqual({0: 0, 1: 215, 2: 300}, max_memory)


Review comments: body: ```suggestion
    offset = 2 if "cpu" in (list(max_memory.keys())) else 1
    for i in range(len(max_memory) - offset):
```
I think this way it seems that the slow tests of `accelerate` are not broken and we also keep the `transformers` tests pass! Can you confirm this pass the `transformers` tests that were broken @ArthurZucker ? 🙏 
diff_hunk: @@ -409,7 +409,7 @@ def get_balanced_memory(
 
     max_memory = get_max_memory(max_memory)
     # The last device is left with max_memory just in case the buffer is not enough.
-    for i in range(len(max_memory) - 1):
+    for i in range(num_devices - 1):
body: Will have a look😄 I am also adding the following 
```python
max_memory = get_balanced_memory(model, max_memory={0: 300, 1: 500, "cpu":1000})
```
diff_hunk: @@ -409,7 +409,7 @@ def get_balanced_memory(
 
     max_memory = get_max_memory(max_memory)
     # The last device is left with max_memory just in case the buffer is not enough.
-    for i in range(len(max_memory) - 1):
+    for i in range(num_devices - 1):
body: Ahhh great! I was missing this test! Thanks for clarifying 💪 
diff_hunk: @@ -409,7 +409,7 @@ def get_balanced_memory(
 
     max_memory = get_max_memory(max_memory)
     # The last device is left with max_memory just in case the buffer is not enough.
-    for i in range(len(max_memory) - 1):
+    for i in range(num_devices - 1):
body: This won't work when the last device is GPU 3 but the user forced 0 at GPU 0, since `num_devices` will be 3 (instead of 4) and thus the max memory will be at GPU 2 instead of 3. (you can try running the multi-GPU tests and they will fail).

Agreed that `len(max_memory)` is not good either, so we should just compute the biggest integer in the keys of `max_memory`:
```py
last_gpu = max(i for i in max_memory if isintance(i, int))
```
then use this.


diff_hunk: @@ -409,7 +409,7 @@ def get_balanced_memory(
 
     max_memory = get_max_memory(max_memory)
     # The last device is left with max_memory just in case the buffer is not enough.
-    for i in range(len(max_memory) - 1):
+    for i in range(num_devices - 1):
body: Okay will add this suggestion 💯 
diff_hunk: @@ -409,7 +409,7 @@ def get_balanced_memory(
 
     max_memory = get_max_memory(max_memory)
     # The last device is left with max_memory just in case the buffer is not enough.
-    for i in range(len(max_memory) - 1):
+    for i in range(num_devices - 1):
body: Oh and we should avoid a GPU that the user set at 0, so actually
```
last_gpu = max(i for i in max_memory if isintance(i, int) and max_memory[i] > 0)
```
diff_hunk: @@ -409,7 +409,7 @@ def get_balanced_memory(
 
     max_memory = get_max_memory(max_memory)
     # The last device is left with max_memory just in case the buffer is not enough.
-    for i in range(len(max_memory) - 1):
+    for i in range(num_devices - 1):

Discussion: body: _The documentation is not available anymore as the PR was closed or merged._