Update cross_entropy_loss.py

This commit is contained in:
Daniel Han 2024-09-25 00:00:51 -07:00
commit b30acca3b6

View file

@ -377,17 +377,17 @@ def fast_cross_entropy_loss(
pass
from transformers.models.llama.modeling_llama import LlamaForCausalLM
from transformers.models.llama.modeling_llama import (
LlamaForCausalLM,
CausalLMOutputWithPast,
Optional,
Union,
Cache,
List,
Tuple,
)
def patch_llama_for_causal_lm():
import transformers.models.llama.modeling_llama
from transformers.models.llama.modeling_llama import (
CausalLMOutputWithPast,
Optional,
Union,
Cache,
List,
Tuple,
)
import inspect, re
function = inspect.getsource(transformers.models.llama.modeling_llama.LlamaForCausalLM.forward)
function = function.split("\n")