fix unittests

Signed-off-by: Oleksii Kuchaiev <okuchaiev@nvidia.com>
NVIDIA · Feb 11, 2020 · 51120d9 · 51120d9
1 parent 781ebc5
commit 51120d9
Show file tree

Hide file tree

Showing 4 changed files with 4 additions and 4 deletions.
diff --git a/nemo/backends/pytorch/common/rnn.py b/nemo/backends/pytorch/common/rnn.py
@@ -70,7 +70,7 @@ def input_ports(self):
         """
         return {
             # 'targets': NeuralType({0: AxisType(BatchTag), 1: AxisType(TimeTag)}),
-            'targets': NeuralType(ChannelType(), ('B', 'T')),
+            'targets': NeuralType(LabelsType(), ('B', 'T')),
             # 'encoder_outputs': NeuralType(
             #   {0: AxisType(BatchTag), 1: AxisType(TimeTag), 2: AxisType(ChannelTag),}, optional=True,
             # ),

diff --git a/nemo/collections/asr/data_layer.py b/nemo/collections/asr/data_layer.py
@@ -102,7 +102,7 @@ def output_ports(self):
             # 'transcript_length': NeuralType({0: AxisType(BatchTag)}),
             'audio_signal': NeuralType(AudioSignal(freq=self._sample_rate), ('B', 'T')),
             'a_sig_length': NeuralType(LengthsType(), tuple('B')),
-            'transcripts': NeuralType(ChannelType(), ('B', 'T')),
+            'transcripts': NeuralType(LabelsType(), ('B', 'T')),
             'transcript_length': NeuralType(LengthsType(), tuple('B')),
         }
 

diff --git a/nemo/collections/asr/losses.py b/nemo/collections/asr/losses.py
@@ -25,7 +25,7 @@ def input_ports(self):
             # "input_length": NeuralType({0: AxisType(BatchTag)}),
             # "target_length": NeuralType({0: AxisType(BatchTag)}),
             "log_probs": NeuralType(LogprobsType(), ('B', 'T', 'D')),
-            "targets": NeuralType(PredictionsType(), ('B', 'T')),
+            "targets": NeuralType(LabelsType(), ('B', 'T')),
             "input_length": NeuralType(LengthsType(), tuple('B')),
             "target_length": NeuralType(LengthsType(), tuple('B')),
         }

diff --git a/tests/asr/test_zeroDS.py b/tests/asr/test_zeroDS.py
@@ -112,7 +112,7 @@ def test_asr_with_zero_ds(self):
                     (AxisType(AxisKind.Batch), AxisType(AxisKind.Dimension, 64), AxisType(AxisKind.Time, 64)),
                 ),
                 "processed_length": NeuralType(LengthsType(), tuple('B')),
-                "transcript": NeuralType(ChannelType(), (AxisType(AxisKind.Batch), AxisType(AxisKind.Time, 64))),
+                "transcript": NeuralType(LabelsType(), (AxisType(AxisKind.Batch), AxisType(AxisKind.Time, 64))),
                 "transcript_length": NeuralType(LengthsType(), tuple('B')),
             },
         )