cyrilvallez HF staff commited on
Commit
c8844c6
·
verified ·
1 Parent(s): 07c2078

Update README.md

Browse files

Correctly set device_map

Files changed (1) hide show
  1. README.md +1 -0
README.md CHANGED
@@ -239,6 +239,7 @@ def split_model(model_name):
239
  device_map['language_model.model.embed_tokens'] = 0
240
  device_map['language_model.output'] = 0
241
  device_map['language_model.model.norm'] = 0
 
242
  device_map['language_model.lm_head'] = 0
243
  device_map[f'language_model.model.layers.{num_layers - 1}'] = 0
244
 
 
239
  device_map['language_model.model.embed_tokens'] = 0
240
  device_map['language_model.output'] = 0
241
  device_map['language_model.model.norm'] = 0
242
+ device_map['language_model.model.rotary_emb'] = 0
243
  device_map['language_model.lm_head'] = 0
244
  device_map[f'language_model.model.layers.{num_layers - 1}'] = 0
245