Input generation fix

modelling_RW.py

@@ -490,7 +490,9 @@ class RWPreTrainedModel(PreTrainedModel):
         # value: [batch_size * num_heads, seq_length, head_dim] -> [batch_size, num_heads, seq_length, head_dim]
         return tuple(
             (
-                layer_past[0].view(batch_size, num_heads, head_dim, seq_length),
+                # layer_past[0].view(batch_size, num_heads, head_dim, seq_length),
+                # layer_past[1].view(batch_size, num_heads, seq_length, head_dim),
+                layer_past[0].view(batch_size, num_heads, seq_length, head_dim),
                 layer_past[1].view(batch_size, num_heads, seq_length, head_dim),
             )
             for layer_past in past_key_value