From a7f87b636ba136287e60b450d22d452741aefcf3 Mon Sep 17 00:00:00 2001
From: Max Ryabinin <mryabinin0@gmail.com>
Date: Sun, 3 Sep 2023 00:49:23 +0300
Subject: [PATCH] Disable the optimization

---
 src/petals/models/llama/block.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/petals/models/llama/block.py b/src/petals/models/llama/block.py
index b7616a1..46898b9 100644
--- 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,