Disable the optimization

8 months ago · a7f87b636b
parent b2ab84cc33
commit a7f87b636b
1 changed files with 1 additions and 1 deletions
--- a/src/petals/models/llama/block.py
+++ b/src/petals/models/llama/block.py
@ -114,7 +114,7 @@ class OptimizedLlamaDecoderLayer(LlamaDecoderLayer):
        self.post_attention_layernorm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)


-class WrappedLlamaBlock(OptimizedLlamaDecoderLayer):
+class WrappedLlamaBlock(LlamaDecoderLayer):
    def forward(
        self,
        hidden_states: torch.Tensor,