diff --git a/src/accelerate/big_modeling.py b/src/accelerate/big_modeling.py
index 83c6738af..e277e89a2 100644
--- a/src/accelerate/big_modeling.py
+++ b/src/accelerate/big_modeling.py
@@ -245,11 +245,15 @@ def dispatch_model(
     check_device_map(model, device_map)
 
     if main_device is None:
-        main_device = [d for d in device_map.values() if d not in ["cpu", "disk"]][0]
+        if set(device_map.values()) == {"cpu"} or set(device_map.values()) == {"cpu", "disk"}:
+            main_device = "cpu"
+        else:
+            main_device = [d for d in device_map.values() if d not in ["cpu", "disk"]][0]
 
-    cpu_modules = [name for name, device in device_map.items() if device == "cpu"]
-    if state_dict is None and len(cpu_modules) > 0:
-        state_dict = extract_submodules_state_dict(model.state_dict(), cpu_modules)
+    if main_device != "cpu":
+        cpu_modules = [name for name, device in device_map.items() if device == "cpu"]
+        if state_dict is None and len(cpu_modules) > 0:
+            state_dict = extract_submodules_state_dict(model.state_dict(), cpu_modules)
 
     disk_modules = [name for name, device in device_map.items() if device == "disk"]
     if offload_dir is None and len(disk_modules) > 0:
@@ -266,7 +270,8 @@ def dispatch_model(
     execution_device = {
         name: main_device if device in ["cpu", "disk"] else device for name, device in device_map.items()
     }
-    offload = {name: device in ["cpu", "disk"] for name, device in device_map.items()}
+    offloaded_devices = ["disk"] if main_device == "cpu" else ["cpu", "disk"]
+    offload = {name: device in offloaded_devices for name, device in device_map.items()}
     save_folder = offload_dir if len(disk_modules) > 0 else None
     if state_dict is not None or save_folder is not None:
         weights_map = OffloadedWeightsLoader(state_dict=state_dict, save_folder=save_folder)



Discussion: body: _The documentation is not available anymore as the PR was closed or merged._
body: Is this expected to work on the `mps` device too? I've checked #596, but the only device I see here is `cpu`:
https://github.com/huggingface/accelerate/blob/0de16441266ba9a799aa338aa75c73bf078429e9/src/accelerate/utils/modeling.py#L464

Furthermore, `get_max_memory` doesn't seem to take it into account: https://github.com/huggingface/accelerate/blob/b816e258a95904736c8e3a5a2b14a18d5f407f4a/src/accelerate/utils/modeling.py#L275-L284

Alternatively, is it possible to override using a custom `device_map`?

Reference: https://github.com/huggingface/diffusers/pull/1042
body: No we don't support the "mps" device at all in big model inference. `device_map="auto"` will load the model on the CPU and then the user can't move them to the MPS device if they want.
body: Ok understood, thanks! However, we get errors after moving to the MPS device, if we use `device_map="auto"`. I'm not sure what's causing it, but some tensors are being moved to the CPU when performing inference. I'll investigate further.
body: That's the side effect of using `device_map="auto"` @Narsil was warning you about. It does more than just load the model with empty weights but adds hooks to the model.
body: Yeah I saw the hooks while stepping with the debugger. One of the reasons why we want to use this is for easier communication (just use `accelerate` and `device_map="auto"`). We'll see what we can do :)
body: @pcuenca Loading weights efficiently is like 2 functions from `accelerate`.

https://github.com/huggingface/accelerate/blob/main/src/accelerate/big_modeling.py#L35-L110

And loading after the fact https://github.com/huggingface/transformers_bloom_parallel/blob/main/generate.py#L211-L270
(Couldn't fetch rapidly from `accelerate` but it has some more complete and robust function).

To be completely honest, I feel like all this should be in Pytorch proper. the `meta` device is really a bad workaround.

In the meantime maybe instead of pushing towards `device_map="auto"` you could instead use the 2 accelerate functions and wrapping them up nicely in `from_pretrained` ?
body: > In the meantime maybe instead of pushing towards `device_map="auto"` you could instead use the 2 accelerate functions and wrapping them up nicely in `from_pretrained` ?

That's what I was thinking, yes! What I'm not sure about is whether to apply those only for the `mps` device (and keep using `device_map` in the others), or something else. I'll do some experiments to see what it looks like in practice. Thanks a lot!