Update llama.py

This commit is contained in:
Daniel Han 2025-02-12 03:03:43 -08:00
commit 648fd043c9

View file

@ -1154,6 +1154,7 @@ def CausalLM_fast_forward(fast_forward_inference):
output = (logits,) + outputs[1:]
return (loss,) + output if loss is not None else output
print("========== dtype = ", logits.dtype)
return CausalLMOutputWithPast(
loss=loss,
logits=logits,