Fix example code on README.md (#21)

- Fix example code on README.md (c4fe18c4102abe2f1b6720f60ea27e1d51a74b3a)


Co-authored-by: Siri Thammarerkrit <[email protected]>

README.md

@@ -158,6 +158,7 @@ def split_model():
  device_map['language_model.output'] = 0
  device_map['language_model.model.norm'] = 0
  device_map['language_model.lm_head'] = 0
+ device_map['language_model.model.rotary_emb'] = 0
  device_map[f'language_model.model.layers.{num_layers - 1}'] = 0
 
  return device_map
@@ -205,6 +206,7 @@ def split_model():
  device_map['language_model.output'] = 0
  device_map['language_model.model.norm'] = 0
  device_map['language_model.lm_head'] = 0
+ device_map['language_model.model.rotary_emb'] = 0
  device_map[f'language_model.model.layers.{num_layers - 1}'] = 0
 
  return device_map