Merge branch 'lgrigoryan/rm-redundant-calculations' of https://github…

….com/lilithgrigoryan/NeMo into lgrigoryan/rm-redundant-calculations
lilithgrigoryan · Oct 11, 2024 · 141ca40 · 141ca40
2 parents 0917ceb + f3a10c4
commit 141ca40
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/nemo/collections/asr/modules/rnnt.py b/nemo/collections/asr/modules/rnnt.py
@@ -1037,7 +1037,11 @@ 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]) if not isinstance(batch_states[sample_id][state_id], torch.Tensor) else 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)