Update llama.py

This commit is contained in:
Daniel Han-Chen 2024-02-14 23:15:45 +11:00
commit 227c26ca44

View file

@ -781,7 +781,7 @@ pass
# Inferene can now be CUDAGraphed, but we shall retain the old rotary embeddings.
# https://github.com/huggingface/transformers/pull/27931
# https://github.com/huggingface/transformers/blob/v4.37.2/src/transformers/models/llama/modeling_llama.py
class LlamaRotaryEmbedding(nn.Module):
class LlamaRotaryEmbedding(torch.nn.Module):
def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None):
super().__init__()