iristun commited on
Commit
c4fe18c
1 Parent(s): 32d8bf4

Fix example code on README.md

Browse files

Earlier, 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.

Files changed (1) hide show
  1. README.md +2 -0
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