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

Update llama.py

Browse source
This commit is contained in:
Daniel Han 2025-03-04 18:18:02 -08:00
commit 6f0857ba46
1 changed files with 3 additions and 3 deletions
Download patch file Download diff file Expand all files Collapse all files

6
unsloth/models/llama.py
View file

@ -1534,11 +1534,11 @@ class LongRopeRotaryEmbedding(torch.nn.Module):
pass
def _wrap_fast_inference(generate):
def _wrap_fast_inference(generate_function):
# Wraps inference with bfloat16 / float16
@torch.inference_mode
def _fast_generate(self, *args, **kwargs):
f"""{getattr(generate, '__doc__', 'Unsloth fast generation')}"""
f"""{getattr(generate_function, '__doc__', 'Unsloth fast generation')}"""
FastLlamaModel.for_inference(self)
@ -1576,7 +1576,7 @@ def _wrap_fast_inference(generate):
# Mixed precision autocast
with torch.autocast(device_type = "cuda", dtype = dtype):
output = generate(*args, **kwargs)
output = generate_function(self, *args, **kwargs)
pass
# Return accelerate back