Cannot batch inference with WavLM in `torchaudio.pipelines`

🐛 Describe the bug

Batch inference with WavLM triggers AssertionError in WavLMSelfAttention module.

import torchaudio
wavlm=torchaudio.pipelines.WAVLM_LARGE.get_model().cuda()
wavlm.extract_features(torch.randn(2,16000,device='cuda'),lengths=torch.tensor([2000,3000],device='cuda'),num_layers=1)

Log:

AssertionError                            Traceback (most recent call last)
[<ipython-input-43-11ed28e9b1d6>](https://localhost:8080/#) in <cell line: 3>()
      1 import torchaudio
      2 wavlm=torchaudio.pipelines.WAVLM_LARGE.get_model().cuda()
----> 3 wavlm.extract_features(torch.randn(2,16000,device='cuda'),lengths=torch.tensor([2000,3000],device='cuda'),num_layers=1)

8 frames
[/usr/local/lib/python3.10/dist-packages/torchaudio/models/wav2vec2/model.py](https://localhost:8080/#) in extract_features(self, waveforms, lengths, num_layers)
     82         """
     83         x, lengths = self.feature_extractor(waveforms, lengths)
---> 84         x = self.encoder.extract_features(x, lengths, num_layers)
     85         return x, lengths
     86 

[/usr/local/lib/python3.10/dist-packages/torchaudio/models/wav2vec2/components.py](https://localhost:8080/#) in extract_features(self, features, lengths, num_layers)
    508     ) -> List[Tensor]:
    509         x, masks = self._preprocess(features, lengths)
--> 510         return self.transformer.get_intermediate_outputs(x, attention_mask=masks, num_layers=num_layers)
    511 
    512 

[/usr/local/lib/python3.10/dist-packages/torchaudio/models/wav2vec2/components.py](https://localhost:8080/#) in get_intermediate_outputs(self, x, attention_mask, num_layers)
    457         x = self._preprocess(x)
    458         for layer in self.layers:
--> 459             x, position_bias = layer(x, attention_mask, position_bias=position_bias)
    460             ret.append(x)
    461             if num_layers is not None and len(ret) >= num_layers:

[/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py](https://localhost:8080/#) in _wrapped_call_impl(self, *args, **kwargs)
   1516             return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
   1517         else:
-> 1518             return self._call_impl(*args, **kwargs)
   1519 
   1520     def _call_impl(self, *args, **kwargs):

[/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py](https://localhost:8080/#) in _call_impl(self, *args, **kwargs)
   1525                 or _global_backward_pre_hooks or _global_backward_hooks
   1526                 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1527             return forward_call(*args, **kwargs)
   1528 
   1529         try:

[/usr/local/lib/python3.10/dist-packages/torchaudio/models/wav2vec2/components.py](https://localhost:8080/#) in forward(self, x, attention_mask, position_bias, key_padding_mask)
    387             x = self.layer_norm(x)
    388 
--> 389         x, position_bias = self.attention(
    390             x, attention_mask=attention_mask, position_bias=position_bias, key_padding_mask=key_padding_mask
    391         )

[/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py](https://localhost:8080/#) in _wrapped_call_impl(self, *args, **kwargs)
   1516             return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
   1517         else:
-> 1518             return self._call_impl(*args, **kwargs)
   1519 
   1520     def _call_impl(self, *args, **kwargs):

[/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py](https://localhost:8080/#) in _call_impl(self, *args, **kwargs)
   1525                 or _global_backward_pre_hooks or _global_backward_hooks
   1526                 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1527             return forward_call(*args, **kwargs)
   1528 
   1529         try:

[/usr/local/lib/python3.10/dist-packages/torchaudio/models/wav2vec2/wavlm_attention.py](https://localhost:8080/#) in forward(self, query, key_padding_mask, attention_mask, position_bias)
    163         bsz, seq_len, embed_dim = query.size()
    164         assert embed_dim == self.embed_dim
--> 165         assert attention_mask is None
    166 
    167         if self.rel_attn_embed is not None and position_bias is None:

AssertionError:

Versions

PyTorch version: 2.1.0+cu118
Is debug build: False
CUDA used to build PyTorch: 11.8
ROCM used to build PyTorch: N/A

OS: Ubuntu 22.04.3 LTS (x86_64)
GCC version: (Ubuntu 11.4.0-1ubuntu1~22.04) 11.4.0
Clang version: 14.0.0-1ubuntu1.1
CMake version: version 3.27.7
Libc version: glibc-2.35

Python version: 3.10.12 (main, Jun 11 2023, 05:26:28) [GCC 11.4.0] (64-bit runtime)
Python platform: Linux-5.15.120+-x86_64-with-glibc2.35
Is CUDA available: True
CUDA runtime version: 11.8.89
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: GPU 0: Tesla T4
Nvidia driver version: 525.105.17
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.8.9.6
/usr/lib/x86_64-linux-gnu/libcudnn_adv_infer.so.8.9.6
/usr/lib/x86_64-linux-gnu/libcudnn_adv_train.so.8.9.6
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_infer.so.8.9.6
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_train.so.8.9.6
/usr/lib/x86_64-linux-gnu/libcudnn_ops_infer.so.8.9.6
/usr/lib/x86_64-linux-gnu/libcudnn_ops_train.so.8.9.6
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture: x86_64
CPU op-mode(s): 32-bit, 64-bit
Address sizes: 46 bits physical, 48 bits virtual
Byte Order: Little Endian
CPU(s): 2
On-line CPU(s) list: 0,1
Vendor ID: GenuineIntel
Model name: Intel(R) Xeon(R) CPU @ 2.00GHz
CPU family: 6
Model: 85
Thread(s) per core: 2
Core(s) per socket: 1
Socket(s): 1
Stepping: 3
BogoMIPS: 4000.42
Flags: fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ss ht syscall nx pdpe1gb rdtscp lm constant_tsc rep_good nopl xtopology nonstop_tsc cpuid tsc_known_freq pni pclmulqdq ssse3 fma cx16 pcid sse4_1 sse4_2 x2apic movbe popcnt aes xsave avx f16c rdrand hypervisor lahf_lm abm 3dnowprefetch invpcid_single ssbd ibrs ibpb stibp fsgsbase tsc_adjust bmi1 hle avx2 smep bmi2 erms invpcid rtm mpx avx512f avx512dq rdseed adx smap clflushopt clwb avx512cd avx512bw avx512vl xsaveopt xsavec xgetbv1 xsaves arat md_clear arch_capabilities
Hypervisor vendor: KVM
Virtualization type: full
L1d cache: 32 KiB (1 instance)
L1i cache: 32 KiB (1 instance)
L2 cache: 1 MiB (1 instance)
L3 cache: 38.5 MiB (1 instance)
NUMA node(s): 1
NUMA node0 CPU(s): 0,1
Vulnerability Itlb multihit: Not affected
Vulnerability L1tf: Mitigation; PTE Inversion
Vulnerability Mds: Vulnerable; SMT Host state unknown
Vulnerability Meltdown: Vulnerable
Vulnerability Mmio stale data: Vulnerable
Vulnerability Retbleed: Vulnerable
Vulnerability Spec store bypass: Vulnerable
Vulnerability Spectre v1: Vulnerable: __user pointer sanitization and usercopy barriers only; no swapgs barriers
Vulnerability Spectre v2: Vulnerable, IBPB: disabled, STIBP: disabled, PBRSB-eIBRS: Not affected
Vulnerability Srbds: Not affected
Vulnerability Tsx async abort: Vulnerable

Versions of relevant libraries:
[pip3] numpy==1.23.5
[pip3] torch==2.1.0+cu118
[pip3] torchaudio==2.1.0+cu118
[pip3] torchdata==0.7.0
[pip3] torchsummary==1.5.1
[pip3] torchtext==0.16.0
[pip3] torchvision==0.16.0+cu118
[pip3] triton==2.1.0
[conda] Could not collect

Start by reproducing the example with WavLM batch inputs, then read torchaudio/models/wav2vec2/wavlm_attention.py around the assertion and follow the call path through components.py. Done means the reported batched extract_features call no longer raises AssertionError, with coverage for the lengths and attention-mask case.