Skip to content

Commit

Permalink
remove torch compile
Browse files Browse the repository at this point in the history
  • Loading branch information
eliebak committed Aug 27, 2024
1 parent 03d67f2 commit 3a45a34
Showing 1 changed file with 1 addition and 1 deletion.
2 changes: 1 addition & 1 deletion src/nanotron/models/llama.py
Original file line number Diff line number Diff line change
Expand Up @@ -165,7 +165,7 @@ def __init__(
bias=False,
async_communication=tp_linear_async_communication and tp_mode is TensorParallelLinearMode.REDUCE_SCATTER,
)
self.split_silu_mul = torch.compile(GLUActivation(config.hidden_act))
self.split_silu_mul = GLUActivation(config.hidden_act)

def forward(self, hidden_states): # [seq_length, batch_size, hidden_dim]
merged_states = self.gate_up_proj(hidden_states)
Expand Down

0 comments on commit 3a45a34

Please sign in to comment.