Repository: huggingface/accelerate

Issue:
Issue with using `dispatch_model` with timm models
### System Info

```Shell
- `Accelerate` version: 0.16.0
- Platform: Linux-5.4.0-1063-aws-x86_64-with-glibc2.27
- Python version: 3.10.8
- Numpy version: 1.24.1
- PyTorch version (GPU?): 1.13.1+cu117 (True)
- `Accelerate` default config:
	Not found
```


### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [ ] One of the scripts in the examples/ folder of Accelerate or an officially supported `no_trainer` script in the `examples` folder of the `transformers` repo (such as `run_no_trainer_glue.py`)
- [X] My own task or dataset (give details below)

### Reproduction

I'm trying to use `accelerate.dispatch_model` on timm's vit and it's not working properly.
When executing: 
```
import accelerate
import timm
import torch

model = timm.create_model("vit_huge_patch14_224")
devices = accelerate.infer_auto_device_map(model, {i: "2GB" for i in range(8)}, no_split_module_classes="Block")
model = accelerate.dispatch_model(model, devices)
images = torch.randn((10, 3, 224, 224)).to(0)
print(model(images))
```
I get the following error:
```
╭───────────────────── Traceback (most recent call last) ──────────────────────╮
│ example.py:9 in <module>              │
│                                                                              │
│    6 devices = accelerate.infer_auto_device_map(model, {i: "2GB" for i in ra │
│    7 model = accelerate.dispatch_model(model, devices)                       │
│    8 images = torch.randn((10, 3, 224, 224)).to(0)                           │
│ ❱  9 print(model(images))                                                    │
│   10                                                                         │
│                                                                              │
│ /home/ubuntu/anaconda3/envs/test-env/lib/python3.10/site-package │
│ s/torch/nn/modules/module.py:1194 in _call_impl                              │
│                                                                              │
│   1191 │   │   # this function, and just call forward.                       │
│   1192 │   │   if not (self._backward_hooks or self._forward_hooks or self._ │
│   1193 │   │   │   │   or _global_forward_hooks or _global_forward_pre_hooks │
│ ❱ 1194 │   │   │   return forward_call(*input, **kwargs)                     │
│   1195 │   │   # Do not call functions when jit is used                      │
│   1196 │   │   full_backward_hooks, non_full_backward_hooks = [], []         │
│   1197 │   │   if self._backward_hooks or _global_backward_hooks:            │
│                                                                              │
│ /home/ubuntu/anaconda3/envs/test-env/lib/python3.10/site-package │
│ s/accelerate/hooks.py:158 in new_forward                                     │
│                                                                              │
│   155 │   │   │   with torch.no_grad():                                      │
│   156 │   │   │   │   output = old_forward(*args, **kwargs)                  │
│   157 │   │   else:                                                          │
│ ❱ 158 │   │   │   output = old_forward(*args, **kwargs)                      │
│   159 │   │   return module._hf_hook.post_forward(module, output)            │
│   160 │                                                                      │
│   161 │   module.forward = new_forward                                       │
│                                                                              │
│ /home/ubuntu/anaconda3/envs/test-env/lib/python3.10/site-package │
│ s/timm/models/vision_transformer.py:549 in forward                           │
│                                                                              │
│    546 │   │   return x if pre_logits else self.head(x)                      │
│    547 │                                                                     │
│    548 │   def forward(self, x):                                             │
│ ❱  549 │   │   x = self.forward_features(x)                                  │
│    550 │   │   x = self.forward_head(x)                                      │
│    551 │   │   return x                                                      │
│    552                                                                       │
│                                                                              │
│ /home/ubuntu/anaconda3/envs/test-env/lib/python3.10/site-package │
│ s/timm/models/vision_transformer.py:533 in forward_features                  │
│                                                                              │
│    530 │                                                                     │
│    531 │   def forward_features(self, x):                                    │
│    532 │   │   x = self.patch_embed(x)                                       │
│ ❱  533 │   │   x = self._pos_embed(x)                                        │
│    534 │   │   x = self.norm_pre(x)                                          │
│    535 │   │   if self.grad_checkpointing and not torch.jit.is_scripting():  │
│    536 │   │   │   x = checkpoint_seq(self.blocks, x)                        │
│                                                                              │
│ /home/ubuntu/anaconda3/envs/test-env/lib/python3.10/site-package │
│ s/timm/models/vision_transformer.py:527 in _pos_embed                        │
│                                                                              │
│    524 │   │   │   # original timm, JAX, and deit vit impl                   │
│    525 │   │   │   # pos_embed has entry for class token, concat then add    │
│    526 │   │   │   if self.cls_token is not None:                            │
│ ❱  527 │   │   │   │   x = torch.cat((self.cls_token.expand(x.shape[0], -1,  │
│    528 │   │   │   x = x + self.pos_embed                                    │
│    529 │   │   return self.pos_drop(x)                                       │
│    530                                                                       │
╰──────────────────────────────────────────────────────────────────────────────╯
RuntimeError: Expected all tensors to be on the same device, but found at least 
two devices, cpu and cuda:0! (when checking argument for argument tensors in 
method wrapper_cat)
```
If I comment out line 533: ` x = self._pos_embed(x)` it runs just fine, so I'm assuming there's and issue when handling `nn.Parameter` inside the dispatched model. Note that the model is moved to the gpu.
I'm also attaching the `device_map` below:
`{'cls_token': 0, 'pos_embed': 0, 'patch_embed': 0, 'pos_drop': 0, 'norm_pre': 0, 'blocks.0': 0, 'blocks.1': 0, 'blocks.2': 0, 'blocks.3': 0, 'blocks.4': 0, 'blocks.5': 0, 'blocks.6': 0, 'blocks.7': 0, 'blocks.8': 0, 'blocks.9': 0, 'blocks.10': 0, 'blocks.11': 1, 'blocks.12': 1, 'blocks.13': 1, 'blocks.14': 1, 'blocks.15': 1, 'blocks.16': 1, 'blocks.17': 1, 'blocks.18': 1, 'blocks.19': 1, 'blocks.20': 1, 'blocks.21': 1, 'blocks.22': 1, 'blocks.23': 2, 'blocks.24': 2, 'blocks.25': 2, 'blocks.26': 2, 'blocks.27': 2, 'blocks.28': 2, 'blocks.29': 2, 'blocks.30': 2, 'blocks.31': 2, 'norm': 2, 'fc_norm': 2, 'head': 2}`


### Expected behavior

```Shell
The model should be split among GPUs and allow forwarding data.
```



Discussion:

Resolving patch:
diff --git a/src/accelerate/big_modeling.py b/src/accelerate/big_modeling.py
--- a/src/accelerate/big_modeling.py
+++ b/src/accelerate/big_modeling.py
@@ -303,6 +303,7 @@ def dispatch_model(
     execution_device = {
         name: main_device if device in ["cpu", "disk"] else device for name, device in device_map.items()
     }
+    execution_device[""] = main_device
     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
diff --git a/src/accelerate/hooks.py b/src/accelerate/hooks.py
--- a/src/accelerate/hooks.py
+++ b/src/accelerate/hooks.py
@@ -454,7 +454,7 @@ def attach_align_device_hook_on_blocks(
     if not isinstance(offload, Mapping):
         offload = {key: offload for key in execution_device.keys()}
 
-    if module_name in execution_device and not offload[module_name]:
+    if module_name in execution_device and module_name in offload and not offload[module_name]:
         hook = AlignDevicesHook(
             execution_device=execution_device[module_name],
             offload_buffers=offload_buffers,
@@ -463,7 +463,7 @@ def attach_align_device_hook_on_blocks(
         )
         add_hook_to_module(module, hook)
         attach_execution_device_hook(module, execution_device[module_name])
-    elif module_name in execution_device:
+    elif module_name in execution_device and module_name in offload:
         attach_align_device_hook(
             module,
             execution_device=execution_device[module_name],
@@ -480,7 +480,7 @@ def attach_align_device_hook_on_blocks(
             module, execution_device[module_name], preload_module_classes=preload_module_classes
         )
     elif module_name == "":
-        hook = AlignDevicesHook(io_same_device=True)
+        hook = AlignDevicesHook(execution_device=execution_device.get(""), io_same_device=True)
         add_hook_to_module(module, hook)
 
     for child_name, child in module.named_children():


Test patch: