Fix example code on README.md
Browse filesEarlier, I got this error message from the example code:
```
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 mat2 in method wrapper_CUDA_bmm)
```
I discovered that the problem originated from the `rotary_emb`.
To resolve this, I added `device_map['language_model.model.rotary_emb'] = 0` to the device map, and it successfully resolved the issue.
README.md
CHANGED
|
@@ -158,6 +158,7 @@ def split_model():
|
|
| 158 |
device_map['language_model.output'] = 0
|
| 159 |
device_map['language_model.model.norm'] = 0
|
| 160 |
device_map['language_model.lm_head'] = 0
|
|
|
|
| 161 |
device_map[f'language_model.model.layers.{num_layers - 1}'] = 0
|
| 162 |
|
| 163 |
return device_map
|
|
@@ -205,6 +206,7 @@ def split_model():
|
|
| 205 |
device_map['language_model.output'] = 0
|
| 206 |
device_map['language_model.model.norm'] = 0
|
| 207 |
device_map['language_model.lm_head'] = 0
|
|
|
|
| 208 |
device_map[f'language_model.model.layers.{num_layers - 1}'] = 0
|
| 209 |
|
| 210 |
return device_map
|
|
|
|
| 158 |
device_map['language_model.output'] = 0
|
| 159 |
device_map['language_model.model.norm'] = 0
|
| 160 |
device_map['language_model.lm_head'] = 0
|
| 161 |
+
device_map['language_model.model.rotary_emb'] = 0
|
| 162 |
device_map[f'language_model.model.layers.{num_layers - 1}'] = 0
|
| 163 |
|
| 164 |
return device_map
|
|
|
|
| 206 |
device_map['language_model.output'] = 0
|
| 207 |
device_map['language_model.model.norm'] = 0
|
| 208 |
device_map['language_model.lm_head'] = 0
|
| 209 |
+
device_map['language_model.model.rotary_emb'] = 0
|
| 210 |
device_map[f'language_model.model.layers.{num_layers - 1}'] = 0
|
| 211 |
|
| 212 |
return device_map
|