Update llama.py

This commit is contained in:
Daniel Han 2024-07-23 11:23:18 -07:00
commit ca3a1b7315

View file

@ -981,8 +981,16 @@ class LlamaRotaryEmbedding(torch.nn.Module):
config = None, # [TODO] Hack to pass in config - need to remove later
):
super().__init__()
if config is not None: return # [TODO] Hack to pass in config - need to remove later
if config is not None:
# [TODO] Hack to pass in config - need to remove later
print(2)
base = config.rope_theta
partial_rotary_factor = config.partial_rotary_factor if hasattr(config, "partial_rotary_factor") else 1.0
dim = int((config.hidden_size // config.num_attention_heads))
device = "cuda"
max_position_embeddings = config.max_position_embeddings
pass
self.dim = dim
self.max_position_embeddings = max_position_embeddings
self.base = base