Fix Qwen2 LoRA loading

predibase · Mar 20, 2024 · 1221ab3 · 1221ab3
1 parent f670cfb
commit 1221ab3
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/server/lorax_server/models/flash_qwen2.py b/server/lorax_server/models/flash_qwen2.py
@@ -147,7 +147,7 @@ def adapter_layers(self) -> List[str]:
         return ADAPTER_LAYERS
 
     def get_num_layers_for_type(self, layer_type: str) -> int:
-        return 1 if layer_type == LM_HEAD else len(self.model.transformer.h)
+        return 1 if layer_type == LM_HEAD else len(self.model.model.layers)
 
     def is_row_parallel(self, layer_type: str) -> bool:
         return layer_type in ROW_PARALLEL