ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Update cross_entropy_loss.py

Browse source
This commit is contained in:
Daniel Han 2024-10-28 14:30:06 -07:00
commit ede5d5893a
1 changed files with 5 additions and 3 deletions
Download patch file Download diff file Expand all files Collapse all files

8
unsloth/kernels/cross_entropy_loss.py
View file

@ -349,7 +349,7 @@ class Fast_CrossEntropyLoss(torch.autograd.Function):
pass
@torch._disable_dynamo
# @torch._disable_dynamo
def fast_cross_entropy_loss(
logits,
labels,
@ -472,7 +472,7 @@ def unpatch_llama_for_causal_lm():
pass
@torch._disable_dynamo
# @torch._disable_dynamo
def UnslothForCausalLMLoss(
logits, labels, vocab_size: int, num_items_in_batch: int = None, ignore_index: int = -100, **kwargs
):
@ -495,7 +495,9 @@ def patch_transformers_losses():
import transformers.loss.loss_utils
except:
logger.warning_once("Unsloth: Cannot patch loss functions - update transformers for faster modules!")
return
pass
import transformers.modeling_utils
LOSS_MAPPING = transformers.loss.loss_utils.LOSS_MAPPING
LOSS_MAPPING["ForCausalLM"] = UnslothForCausalLMLoss