fix dosctrings

Signed-off-by: lilithgrigoryan <lgrigoryan@nvidia.com>
NVIDIA · Oct 11, 2024 · 0917ceb · 0917ceb
1 parent 098a1c2
commit 0917ceb
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 6 deletions.
diff --git a/nemo/collections/asr/modules/rnnt.py b/nemo/collections/asr/modules/rnnt.py
@@ -986,15 +986,15 @@ def batch_initialize_states(self, decoder_states: List[List[torch.Tensor]]) -> L
 
         Args:
             decoder_states (list of list of list of torch.Tensor): list of decoder states
-                [B, layer, L, H]
+                [B, C, L, H]
                     - B: Batch size.
-                    - layer: e.g., for LSTM, this is 2: hidden and cell states
+                    - C: e.g., for LSTM, this is 2: hidden and cell states
                     - L: Number of layers in prediction RNN.
                     - H: Dimensionality of the hidden state.
 
         Returns:
             batch_states (list of torch.Tensor): batch of decoder states
-                [layer x torch.Tensor[L x B x H]
+                [C x torch.Tensor[L x B x H]
         """
         # stack decoder states into tensor of shape [B x layers x L x H]
         # permute to the target shape [layers x L x B x H]

diff --git a/nemo/collections/asr/modules/rnnt_abstract.py b/nemo/collections/asr/modules/rnnt_abstract.py
@@ -248,15 +248,15 @@ def batch_initialize_states(self, decoder_states: List[List[torch.Tensor]]):
 
         Args:
             decoder_states (list of list of list of torch.Tensor): list of decoder states
-                [B, layer, L, H]
+                [B, C, L, H]
                     - B: Batch size.
-                    - layer: e.g., for LSTM, this is 2: hidden and cell states
+                    - C: e.g., for LSTM, this is 2: hidden and cell states
                     - L: Number of layers in prediction RNN.
                     - H: Dimensionality of the hidden state.
 
         Returns:
             batch_states (list of torch.Tensor): batch of decoder states
-                [layer x torch.Tensor[L x B x H]
+                [C x torch.Tensor[L x B x H]
         """
         raise NotImplementedError()