Fix HF wav2vec2 test (#1585)

3f0febc4 · moto · GitHub · 7deea259 · 3f0febc4
Unverified Commit 3f0febc4 authored Jun 18, 2021 by moto Committed by GitHub Jun 18, 2021
Hide whitespace changes
Inline Side-by-side

Showing with 1 addition and 1 deletion

test/torchaudio_unittest/models/wav2vec2/huggingface_intergration_test.py ...unittest/models/wav2vec2/huggingface_intergration_test.py +1 -1

No files found.
--- a/test/torchaudio_unittest/models/wav2vec2/huggingface_intergration_test.py
+++ b/test/torchaudio_unittest/models/wav2vec2/huggingface_intergration_test.py
@@ -86,7 +86,7 @@ class TestHFIntegration(TorchaudioTestCase):
        self.assertEqual(ref, hyp)
        # Feature projection
        x = torch.randn(3, 10, config['conv_dim'][-1])
-        ref = original.wav2vec2.feature_projection(x)
+        ref = original.wav2vec2.feature_projection(x)[0]
        hyp = imported.encoder.feature_projection(x)
        self.assertEqual(ref, hyp)
        # Convolutional Positional Encoder