diff --git a/egs/librispeech/ASR/vq_pruned_transducer_stateless2/conformer.py b/egs/librispeech/ASR/vq_pruned_transducer_stateless2/conformer.py
index e71911c35..15c1e89aa 100644
--- a/egs/librispeech/ASR/vq_pruned_transducer_stateless2/conformer.py
+++ b/egs/librispeech/ASR/vq_pruned_transducer_stateless2/conformer.py
@@ -142,7 +142,7 @@ class Conformer(EncoderInterface):
 
         if len(self.output_layers) > 1:
             assert len(self.output_layers) == len(layers_result)
-            return x, lengths, layers_result[0]
+            return x, lengths, layers_result[0].permute(1, 0, 2)
         return x, lengths