Guard model has lm_head

predibase · Nov 1, 2024 · 43c129b · 43c129b
1 parent ad460c0
commit 43c129b
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/server/lorax_server/models/flash_causal_lm.py b/server/lorax_server/models/flash_causal_lm.py
@@ -1627,7 +1627,7 @@ def forward(self, batch: FlashCausalLMBatch, adapter_data: AdapterBatchData) ->
                 lm_head_indices=batch.prefill_head_indices,
             )
 
-            if skip_lm_head:
+            if skip_lm_head and hasattr(self.model, "lm_head"):
                 # re-run through the LM head as the graph did not capture it
                 out = self.model.lm_head(out[0], adapter_data)