Fix device_map

.gitattributes

@@ -36,3 +36,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 examples/red-panda.mp4 filter=lfs diff=lfs merge=lfs -text
 eval_milebench/TextNeedleInAHaystack/TextNeedleInAHaystack_240803235133.json filter=lfs diff=lfs merge=lfs -text
 eval_milebench/TextNeedleInAHaystack/pred.json filter=lfs diff=lfs merge=lfs -text
+examples/image2.jpg filter=lfs diff=lfs merge=lfs -text

README.md

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