fix test

Signed-off-by: lilithgrigoryan <[email protected]>

nemo/collections/asr/modules/rnnt.py

@@ -1037,7 +1037,7 @@ def batch_concat_states(self, batch_states: List[List[torch.Tensor]]) -> List[to
         for state_id in range(len(batch_states[0])):
             batch_list = []
             for sample_id in range(len(batch_states)):
-                tensor = torch.stack(batch_states[sample_id][state_id])  # [L, H]
+                tensor = torch.stack(batch_states[sample_id][state_id]) if not isinstance(batch_states[sample_id][state_id], torch.Tensor) else batch_states[sample_id][state_id]   # [L, H]
                 tensor = tensor.unsqueeze(0)  # [1, L, H]
                 batch_list.append(tensor)