diff --git a/server/text_generation_server/models/rw.py b/server/text_generation_server/models/rw.py
index 92bb135b..d97c1c73 100644
--- a/server/text_generation_server/models/rw.py
+++ b/server/text_generation_server/models/rw.py
@@ -67,18 +67,6 @@ class RW(CausalLM):
         self, input_ids, attention_mask, position_ids, past_key_values: Optional = None
     ) -> Tuple[torch.Tensor, List[Tuple[torch.Tensor, torch.Tensor]]]:
         # Model Forward
-        if past_key_values is not None:
-            reshaped_past_key_values = []
-            for layer in past_key_values:
-                past_keys, past_values = layer
-                reshaped_past_key_values.append(
-                    (
-                        past_keys.view(-1, *past_keys.shape[-2:]),
-                        past_values.view(-1, *past_values.shape[-2:]),
-                    )
-                )
-            past_key_values = reshaped_past_key_values
-
         outputs = self.model.forward(
             input_ids=input_ids,
             attention_mask=attention_mask,