Update llama.py

This commit is contained in:
Daniel Han 2024-07-07 15:46:36 -07:00
commit 55bf35be5d

View file

@ -1037,9 +1037,9 @@ def _wrap_fast_inference(generate, device_type, dtype, model):
kwargs["cache_implementation"] = "dynamic"
# Set pad token
old_pad_token_id = getattr(model.config, "pad_token_id", None)
old_eos_token_id = getattr(model.config, "eos_token_id", None)
model.config.pad_token_id = old_eos_token_id
# old_pad_token_id = getattr(model.config, "pad_token_id", None)
# old_eos_token_id = getattr(model.config, "eos_token_id", None)
# model.config.pad_token_id = old_eos_token_id
# Autocasted
with torch.autocast(device_type = device_type, dtype = dtype):
@ -1047,7 +1047,7 @@ def _wrap_fast_inference(generate, device_type, dtype, model):
pass
# Revert
model.config.pad_token_id = old_pad_token_id
# model.config.pad_token_id = old_pad_token_id
# Unset a flag for generation!
internal_model = model