Dequantize WhisperVQ back to the original audio

[hungphongtrn]

Hi, thank you for the amazing works!

I'm testing the whisper-vq-stoks-medium-en+pl.model with RQBottleneckTransformer inside vq_stoks. I can encode an audio into stokens using encode_audio.

However, when I try to dequantize the stokens using dequantize, I am encountering the error:

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
Cell In[16], [line 1](vscode-notebook-cell:?execution_count=16&line=1)
----> [1](vscode-notebook-cell:?execution_count=16&line=1) vq_model.dequantize(stokens_female[0])

File ~/Workspace/WhisperVQ/.venv/lib/python3.11/site-packages/whisperspeech/vq_stoks.py:517, in RQBottleneckTransformer.dequantize(self, stoks)
    [515](https://vscode-remote+attached-002dcontainer-002b7b22636f6e7461696e65724e616d65223a222f72657365617263685f6c6162227d-0040ssh-002dremote-002b7b22686f73744e616d65223a224a616e2d34784136303030227d.vscode-resource.vscode-cdn.net/home/phong/Workspace/WhisperVQ/~/Workspace/WhisperVQ/.venv/lib/python3.11/site-packages/whisperspeech/vq_stoks.py:515) positions = torch.arange(0, x.shape[-2], dtype=torch.long, device=x.device)
    [516](https://vscode-remote+attached-002dcontainer-002b7b22636f6e7461696e65724e616d65223a222f72657365617263685f6c6162227d-0040ssh-002dremote-002b7b22686f73744e616d65223a224a616e2d34784136303030227d.vscode-resource.vscode-cdn.net/home/phong/Workspace/WhisperVQ/~/Workspace/WhisperVQ/.venv/lib/python3.11/site-packages/whisperspeech/vq_stoks.py:516) x = x + self.positional_embedding(positions)
--> [517](https://vscode-remote+attached-002dcontainer-002b7b22636f6e7461696e65724e616d65223a222f72657365617263685f6c6162227d-0040ssh-002dremote-002b7b22686f73744e616d65223a224a616e2d34784136303030227d.vscode-resource.vscode-cdn.net/home/phong/Workspace/WhisperVQ/~/Workspace/WhisperVQ/.venv/lib/python3.11/site-packages/whisperspeech/vq_stoks.py:517) return self.ln_post(self.out_blocks(x))

File ~/Workspace/WhisperVQ/.venv/lib/python3.11/site-packages/torch/nn/modules/module.py:1553, in Module._wrapped_call_impl(self, *args, **kwargs)
   [1551](https://vscode-remote+attached-002dcontainer-002b7b22636f6e7461696e65724e616d65223a222f72657365617263685f6c6162227d-0040ssh-002dremote-002b7b22686f73744e616d65223a224a616e2d34784136303030227d.vscode-resource.vscode-cdn.net/home/phong/Workspace/WhisperVQ/~/Workspace/WhisperVQ/.venv/lib/python3.11/site-packages/torch/nn/modules/module.py:1551)     return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
   [1552](https://vscode-remote+attached-002dcontainer-002b7b22636f6e7461696e65724e616d65223a222f72657365617263685f6c6162227d-0040ssh-002dremote-002b7b22686f73744e616d65223a224a616e2d34784136303030227d.vscode-resource.vscode-cdn.net/home/phong/Workspace/WhisperVQ/~/Workspace/WhisperVQ/.venv/lib/python3.11/site-packages/torch/nn/modules/module.py:1552) else:
-> [1553](https://vscode-remote+attached-002dcontainer-002b7b22636f6e7461696e65724e616d65223a222f72657365617263685f6c6162227d-0040ssh-002dremote-002b7b22686f73744e616d65223a224a616e2d34784136303030227d.vscode-resource.vscode-cdn.net/home/phong/Workspace/WhisperVQ/~/Workspace/WhisperVQ/.venv/lib/python3.11/site-packages/torch/nn/modules/module.py:1553)     return self._call_impl(*args, **kwargs)

File ~/Workspace/WhisperVQ/.venv/lib/python3.11/site-packages/torch/nn/modules/module.py:1562, in Module._call_impl(self, *args, **kwargs)
   [1557](https://vscode-remote+attached-002dcontainer-002b7b22636f6e7461696e65724e616d65223a222f72657365617263685f6c6162227d-0040ssh-002dremote-002b7b22686f73744e616d65223a224a616e2d34784136303030227d.vscode-resource.vscode-cdn.net/home/phong/Workspace/WhisperVQ/~/Workspace/WhisperVQ/.venv/lib/python3.11/site-packages/torch/nn/modules/module.py:1557) # If we don't have any hooks, we want to skip the rest of the logic in
   [1558](https://vscode-remote+attached-002dcontainer-002b7b22636f6e7461696e65724e616d65223a222f72657365617263685f6c6162227d-0040ssh-002dremote-002b7b22686f73744e616d65223a224a616e2d34784136303030227d.vscode-resource.vscode-cdn.net/home/phong/Workspace/WhisperVQ/~/Workspace/WhisperVQ/.venv/lib/python3.11/site-packages/torch/nn/modules/module.py:1558) # this function, and just call forward.
   [1559](https://vscode-remote+attached-002dcontainer-002b7b22636f6e7461696e65724e616d65223a222f72657365617263685f6c6162227d-0040ssh-002dremote-002b7b22686f73744e616d65223a224a616e2d34784136303030227d.vscode-resource.vscode-cdn.net/home/phong/Workspace/WhisperVQ/~/Workspace/WhisperVQ/.venv/lib/python3.11/site-packages/torch/nn/modules/module.py:1559) if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
   [1560](https://vscode-remote+attached-002dcontainer-002b7b22636f6e7461696e65724e616d65223a222f72657365617263685f6c6162227d-0040ssh-002dremote-002b7b22686f73744e616d65223a224a616e2d34784136303030227d.vscode-resource.vscode-cdn.net/home/phong/Workspace/WhisperVQ/~/Workspace/WhisperVQ/.venv/lib/python3.11/site-packages/torch/nn/modules/module.py:1560)         or _global_backward_pre_hooks or _global_backward_hooks
   [1561](https://vscode-remote+attached-002dcontainer-002b7b22636f6e7461696e65724e616d65223a222f72657365617263685f6c6162227d-0040ssh-002dremote-002b7b22686f73744e616d65223a224a616e2d34784136303030227d.vscode-resource.vscode-cdn.net/home/phong/Workspace/WhisperVQ/~/Workspace/WhisperVQ/.venv/lib/python3.11/site-packages/torch/nn/modules/module.py:1561)         or _global_forward_hooks or _global_forward_pre_hooks):
-> [1562](https://vscode-remote+attached-002dcontainer-002b7b22636f6e7461696e65724e616d65223a222f72657365617263685f6c6162227d-0040ssh-002dremote-002b7b22686f73744e616d65223a224a616e2d34784136303030227d.vscode-resource.vscode-cdn.net/home/phong/Workspace/WhisperVQ/~/Workspace/WhisperVQ/.venv/lib/python3.11/site-packages/torch/nn/modules/module.py:1562)     return forward_call(*args, **kwargs)
   [1564](https://vscode-remote+attached-002dcontainer-002b7b22636f6e7461696e65724e616d65223a222f72657365617263685f6c6162227d-0040ssh-002dremote-002b7b22686f73744e616d65223a224a616e2d34784136303030227d.vscode-resource.vscode-cdn.net/home/phong/Workspace/WhisperVQ/~/Workspace/WhisperVQ/.venv/lib/python3.11/site-packages/torch/nn/modules/module.py:1564) try:
   [1565](https://vscode-remote+attached-002dcontainer-002b7b22636f6e7461696e65724e616d65223a222f72657365617263685f6c6162227d-0040ssh-002dremote-002b7b22686f73744e616d65223a224a616e2d34784136303030227d.vscode-resource.vscode-cdn.net/home/phong/Workspace/WhisperVQ/~/Workspace/WhisperVQ/.venv/lib/python3.11/site-packages/torch/nn/modules/module.py:1565)     result = None

File ~/Workspace/WhisperVQ/.venv/lib/python3.11/site-packages/torch/nn/modules/container.py:219, in Sequential.forward(self, input)
    [217](https://vscode-remote+attached-002dcontainer-002b7b22636f6e7461696e65724e616d65223a222f72657365617263685f6c6162227d-0040ssh-002dremote-002b7b22686f73744e616d65223a224a616e2d34784136303030227d.vscode-resource.vscode-cdn.net/home/phong/Workspace/WhisperVQ/~/Workspace/WhisperVQ/.venv/lib/python3.11/site-packages/torch/nn/modules/container.py:217) def forward(self, input):
    [218](https://vscode-remote+attached-002dcontainer-002b7b22636f6e7461696e65724e616d65223a222f72657365617263685f6c6162227d-0040ssh-002dremote-002b7b22686f73744e616d65223a224a616e2d34784136303030227d.vscode-resource.vscode-cdn.net/home/phong/Workspace/WhisperVQ/~/Workspace/WhisperVQ/.venv/lib/python3.11/site-packages/torch/nn/modules/container.py:218)     for module in self:
--> [219](https://vscode-remote+attached-002dcontainer-002b7b22636f6e7461696e65724e616d65223a222f72657365617263685f6c6162227d-0040ssh-002dremote-002b7b22686f73744e616d65223a224a616e2d34784136303030227d.vscode-resource.vscode-cdn.net/home/phong/Workspace/WhisperVQ/~/Workspace/WhisperVQ/.venv/lib/python3.11/site-packages/torch/nn/modules/container.py:219)         input = module(input)
    [220](https://vscode-remote+attached-002dcontainer-002b7b22636f6e7461696e65724e616d65223a222f72657365617263685f6c6162227d-0040ssh-002dremote-002b7b22686f73744e616d65223a224a616e2d34784136303030227d.vscode-resource.vscode-cdn.net/home/phong/Workspace/WhisperVQ/~/Workspace/WhisperVQ/.venv/lib/python3.11/site-packages/torch/nn/modules/container.py:220)     return input

File ~/Workspace/WhisperVQ/.venv/lib/python3.11/site-packages/torch/nn/modules/module.py:1553, in Module._wrapped_call_impl(self, *args, **kwargs)
   [1551](https://vscode-remote+attached-002dcontainer-002b7b22636f6e7461696e65724e616d65223a222f72657365617263685f6c6162227d-0040ssh-002dremote-002b7b22686f73744e616d65223a224a616e2d34784136303030227d.vscode-resource.vscode-cdn.net/home/phong/Workspace/WhisperVQ/~/Workspace/WhisperVQ/.venv/lib/python3.11/site-packages/torch/nn/modules/module.py:1551)     return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
   [1552](https://vscode-remote+attached-002dcontainer-002b7b22636f6e7461696e65724e616d65223a222f72657365617263685f6c6162227d-0040ssh-002dremote-002b7b22686f73744e616d65223a224a616e2d34784136303030227d.vscode-resource.vscode-cdn.net/home/phong/Workspace/WhisperVQ/~/Workspace/WhisperVQ/.venv/lib/python3.11/site-packages/torch/nn/modules/module.py:1552) else:
-> [1553](https://vscode-remote+attached-002dcontainer-002b7b22636f6e7461696e65724e616d65223a222f72657365617263685f6c6162227d-0040ssh-002dremote-002b7b22686f73744e616d65223a224a616e2d34784136303030227d.vscode-resource.vscode-cdn.net/home/phong/Workspace/WhisperVQ/~/Workspace/WhisperVQ/.venv/lib/python3.11/site-packages/torch/nn/modules/module.py:1553)     return self._call_impl(*args, **kwargs)

File ~/Workspace/WhisperVQ/.venv/lib/python3.11/site-packages/torch/nn/modules/module.py:1562, in Module._call_impl(self, *args, **kwargs)
   [1557](https://vscode-remote+attached-002dcontainer-002b7b22636f6e7461696e65724e616d65223a222f72657365617263685f6c6162227d-0040ssh-002dremote-002b7b22686f73744e616d65223a224a616e2d34784136303030227d.vscode-resource.vscode-cdn.net/home/phong/Workspace/WhisperVQ/~/Workspace/WhisperVQ/.venv/lib/python3.11/site-packages/torch/nn/modules/module.py:1557) # If we don't have any hooks, we want to skip the rest of the logic in
   [1558](https://vscode-remote+attached-002dcontainer-002b7b22636f6e7461696e65724e616d65223a222f72657365617263685f6c6162227d-0040ssh-002dremote-002b7b22686f73744e616d65223a224a616e2d34784136303030227d.vscode-resource.vscode-cdn.net/home/phong/Workspace/WhisperVQ/~/Workspace/WhisperVQ/.venv/lib/python3.11/site-packages/torch/nn/modules/module.py:1558) # this function, and just call forward.
   [1559](https://vscode-remote+attached-002dcontainer-002b7b22636f6e7461696e65724e616d65223a222f72657365617263685f6c6162227d-0040ssh-002dremote-002b7b22686f73744e616d65223a224a616e2d34784136303030227d.vscode-resource.vscode-cdn.net/home/phong/Workspace/WhisperVQ/~/Workspace/WhisperVQ/.venv/lib/python3.11/site-packages/torch/nn/modules/module.py:1559) if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
   [1560](https://vscode-remote+attached-002dcontainer-002b7b22636f6e7461696e65724e616d65223a222f72657365617263685f6c6162227d-0040ssh-002dremote-002b7b22686f73744e616d65223a224a616e2d34784136303030227d.vscode-resource.vscode-cdn.net/home/phong/Workspace/WhisperVQ/~/Workspace/WhisperVQ/.venv/lib/python3.11/site-packages/torch/nn/modules/module.py:1560)         or _global_backward_pre_hooks or _global_backward_hooks
   [1561](https://vscode-remote+attached-002dcontainer-002b7b22636f6e7461696e65724e616d65223a222f72657365617263685f6c6162227d-0040ssh-002dremote-002b7b22686f73744e616d65223a224a616e2d34784136303030227d.vscode-resource.vscode-cdn.net/home/phong/Workspace/WhisperVQ/~/Workspace/WhisperVQ/.venv/lib/python3.11/site-packages/torch/nn/modules/module.py:1561)         or _global_forward_hooks or _global_forward_pre_hooks):
-> [1562](https://vscode-remote+attached-002dcontainer-002b7b22636f6e7461696e65724e616d65223a222f72657365617263685f6c6162227d-0040ssh-002dremote-002b7b22686f73744e616d65223a224a616e2d34784136303030227d.vscode-resource.vscode-cdn.net/home/phong/Workspace/WhisperVQ/~/Workspace/WhisperVQ/.venv/lib/python3.11/site-packages/torch/nn/modules/module.py:1562)     return forward_call(*args, **kwargs)
   [1564](https://vscode-remote+attached-002dcontainer-002b7b22636f6e7461696e65724e616d65223a222f72657365617263685f6c6162227d-0040ssh-002dremote-002b7b22686f73744e616d65223a224a616e2d34784136303030227d.vscode-resource.vscode-cdn.net/home/phong/Workspace/WhisperVQ/~/Workspace/WhisperVQ/.venv/lib/python3.11/site-packages/torch/nn/modules/module.py:1564) try:
   [1565](https://vscode-remote+attached-002dcontainer-002b7b22636f6e7461696e65724e616d65223a222f72657365617263685f6c6162227d-0040ssh-002dremote-002b7b22686f73744e616d65223a224a616e2d34784136303030227d.vscode-resource.vscode-cdn.net/home/phong/Workspace/WhisperVQ/~/Workspace/WhisperVQ/.venv/lib/python3.11/site-packages/torch/nn/modules/module.py:1565)     result = None

File ~/Workspace/WhisperVQ/.venv/lib/python3.11/site-packages/whisperspeech/modules.py:222, in ResidualAttentionBlock.forward(self, x, x_positions, xa, xa_positions, causal, mask)
    [212](https://vscode-remote+attached-002dcontainer-002b7b22636f6e7461696e65724e616d65223a222f72657365617263685f6c6162227d-0040ssh-002dremote-002b7b22686f73744e616d65223a224a616e2d34784136303030227d.vscode-resource.vscode-cdn.net/home/phong/Workspace/WhisperVQ/~/Workspace/WhisperVQ/.venv/lib/python3.11/site-packages/whisperspeech/modules.py:212) def forward(
    [213](https://vscode-remote+attached-002dcontainer-002b7b22636f6e7461696e65724e616d65223a222f72657365617263685f6c6162227d-0040ssh-002dremote-002b7b22686f73744e616d65223a224a616e2d34784136303030227d.vscode-resource.vscode-cdn.net/home/phong/Workspace/WhisperVQ/~/Workspace/WhisperVQ/.venv/lib/python3.11/site-packages/whisperspeech/modules.py:213)     self,
    [214](https://vscode-remote+attached-002dcontainer-002b7b22636f6e7461696e65724e616d65223a222f72657365617263685f6c6162227d-0040ssh-002dremote-002b7b22686f73744e616d65223a224a616e2d34784136303030227d.vscode-resource.vscode-cdn.net/home/phong/Workspace/WhisperVQ/~/Workspace/WhisperVQ/.venv/lib/python3.11/site-packages/whisperspeech/modules.py:214)     x: Tensor,
   (...)
    [219](https://vscode-remote+attached-002dcontainer-002b7b22636f6e7461696e65724e616d65223a222f72657365617263685f6c6162227d-0040ssh-002dremote-002b7b22686f73744e616d65223a224a616e2d34784136303030227d.vscode-resource.vscode-cdn.net/home/phong/Workspace/WhisperVQ/~/Workspace/WhisperVQ/.venv/lib/python3.11/site-packages/whisperspeech/modules.py:219)     mask=None,
    [220](https://vscode-remote+attached-002dcontainer-002b7b22636f6e7461696e65724e616d65223a222f72657365617263685f6c6162227d-0040ssh-002dremote-002b7b22686f73744e616d65223a224a616e2d34784136303030227d.vscode-resource.vscode-cdn.net/home/phong/Workspace/WhisperVQ/~/Workspace/WhisperVQ/.venv/lib/python3.11/site-packages/whisperspeech/modules.py:220) ):
    [221](https://vscode-remote+attached-002dcontainer-002b7b22636f6e7461696e65724e616d65223a222f72657365617263685f6c6162227d-0040ssh-002dremote-002b7b22686f73744e616d65223a224a616e2d34784136303030227d.vscode-resource.vscode-cdn.net/home/phong/Workspace/WhisperVQ/~/Workspace/WhisperVQ/.venv/lib/python3.11/site-packages/whisperspeech/modules.py:221)     lnx = self.attn_ln(x)
--> [222](https://vscode-remote+attached-002dcontainer-002b7b22636f6e7461696e65724e616d65223a222f72657365617263685f6c6162227d-0040ssh-002dremote-002b7b22686f73744e616d65223a224a616e2d34784136303030227d.vscode-resource.vscode-cdn.net/home/phong/Workspace/WhisperVQ/~/Workspace/WhisperVQ/.venv/lib/python3.11/site-packages/whisperspeech/modules.py:222)     x = x + self.attn(lnx, x_positions, lnx, x_positions, causal=causal, mask=mask)
    [223](https://vscode-remote+attached-002dcontainer-002b7b22636f6e7461696e65724e616d65223a222f72657365617263685f6c6162227d-0040ssh-002dremote-002b7b22686f73744e616d65223a224a616e2d34784136303030227d.vscode-resource.vscode-cdn.net/home/phong/Workspace/WhisperVQ/~/Workspace/WhisperVQ/.venv/lib/python3.11/site-packages/whisperspeech/modules.py:223)     if self.cross_attn:
    [224](https://vscode-remote+attached-002dcontainer-002b7b22636f6e7461696e65724e616d65223a222f72657365617263685f6c6162227d-0040ssh-002dremote-002b7b22686f73744e616d65223a224a616e2d34784136303030227d.vscode-resource.vscode-cdn.net/home/phong/Workspace/WhisperVQ/~/Workspace/WhisperVQ/.venv/lib/python3.11/site-packages/whisperspeech/modules.py:224)         lnx = self.cross_attn_ln(x)

File ~/Workspace/WhisperVQ/.venv/lib/python3.11/site-packages/torch/nn/modules/module.py:1553, in Module._wrapped_call_impl(self, *args, **kwargs)
   [1551](https://vscode-remote+attached-002dcontainer-002b7b22636f6e7461696e65724e616d65223a222f72657365617263685f6c6162227d-0040ssh-002dremote-002b7b22686f73744e616d65223a224a616e2d34784136303030227d.vscode-resource.vscode-cdn.net/home/phong/Workspace/WhisperVQ/~/Workspace/WhisperVQ/.venv/lib/python3.11/site-packages/torch/nn/modules/module.py:1551)     return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
   [1552](https://vscode-remote+attached-002dcontainer-002b7b22636f6e7461696e65724e616d65223a222f72657365617263685f6c6162227d-0040ssh-002dremote-002b7b22686f73744e616d65223a224a616e2d34784136303030227d.vscode-resource.vscode-cdn.net/home/phong/Workspace/WhisperVQ/~/Workspace/WhisperVQ/.venv/lib/python3.11/site-packages/torch/nn/modules/module.py:1552) else:
-> [1553](https://vscode-remote+attached-002dcontainer-002b7b22636f6e7461696e65724e616d65223a222f72657365617263685f6c6162227d-0040ssh-002dremote-002b7b22686f73744e616d65223a224a616e2d34784136303030227d.vscode-resource.vscode-cdn.net/home/phong/Workspace/WhisperVQ/~/Workspace/WhisperVQ/.venv/lib/python3.11/site-packages/torch/nn/modules/module.py:1553)     return self._call_impl(*args, **kwargs)

File ~/Workspace/WhisperVQ/.venv/lib/python3.11/site-packages/torch/nn/modules/module.py:1562, in Module._call_impl(self, *args, **kwargs)
   [1557](https://vscode-remote+attached-002dcontainer-002b7b22636f6e7461696e65724e616d65223a222f72657365617263685f6c6162227d-0040ssh-002dremote-002b7b22686f73744e616d65223a224a616e2d34784136303030227d.vscode-resource.vscode-cdn.net/home/phong/Workspace/WhisperVQ/~/Workspace/WhisperVQ/.venv/lib/python3.11/site-packages/torch/nn/modules/module.py:1557) # If we don't have any hooks, we want to skip the rest of the logic in
   [1558](https://vscode-remote+attached-002dcontainer-002b7b22636f6e7461696e65724e616d65223a222f72657365617263685f6c6162227d-0040ssh-002dremote-002b7b22686f73744e616d65223a224a616e2d34784136303030227d.vscode-resource.vscode-cdn.net/home/phong/Workspace/WhisperVQ/~/Workspace/WhisperVQ/.venv/lib/python3.11/site-packages/torch/nn/modules/module.py:1558) # this function, and just call forward.
   [1559](https://vscode-remote+attached-002dcontainer-002b7b22636f6e7461696e65724e616d65223a222f72657365617263685f6c6162227d-0040ssh-002dremote-002b7b22686f73744e616d65223a224a616e2d34784136303030227d.vscode-resource.vscode-cdn.net/home/phong/Workspace/WhisperVQ/~/Workspace/WhisperVQ/.venv/lib/python3.11/site-packages/torch/nn/modules/module.py:1559) if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
   [1560](https://vscode-remote+attached-002dcontainer-002b7b22636f6e7461696e65724e616d65223a222f72657365617263685f6c6162227d-0040ssh-002dremote-002b7b22686f73744e616d65223a224a616e2d34784136303030227d.vscode-resource.vscode-cdn.net/home/phong/Workspace/WhisperVQ/~/Workspace/WhisperVQ/.venv/lib/python3.11/site-packages/torch/nn/modules/module.py:1560)         or _global_backward_pre_hooks or _global_backward_hooks
   [1561](https://vscode-remote+attached-002dcontainer-002b7b22636f6e7461696e65724e616d65223a222f72657365617263685f6c6162227d-0040ssh-002dremote-002b7b22686f73744e616d65223a224a616e2d34784136303030227d.vscode-resource.vscode-cdn.net/home/phong/Workspace/WhisperVQ/~/Workspace/WhisperVQ/.venv/lib/python3.11/site-packages/torch/nn/modules/module.py:1561)         or _global_forward_hooks or _global_forward_pre_hooks):
-> [1562](https://vscode-remote+attached-002dcontainer-002b7b22636f6e7461696e65724e616d65223a222f72657365617263685f6c6162227d-0040ssh-002dremote-002b7b22686f73744e616d65223a224a616e2d34784136303030227d.vscode-resource.vscode-cdn.net/home/phong/Workspace/WhisperVQ/~/Workspace/WhisperVQ/.venv/lib/python3.11/site-packages/torch/nn/modules/module.py:1562)     return forward_call(*args, **kwargs)
   [1564](https://vscode-remote+attached-002dcontainer-002b7b22636f6e7461696e65724e616d65223a222f72657365617263685f6c6162227d-0040ssh-002dremote-002b7b22686f73744e616d65223a224a616e2d34784136303030227d.vscode-resource.vscode-cdn.net/home/phong/Workspace/WhisperVQ/~/Workspace/WhisperVQ/.venv/lib/python3.11/site-packages/torch/nn/modules/module.py:1564) try:
   [1565](https://vscode-remote+attached-002dcontainer-002b7b22636f6e7461696e65724e616d65223a222f72657365617263685f6c6162227d-0040ssh-002dremote-002b7b22686f73744e616d65223a224a616e2d34784136303030227d.vscode-resource.vscode-cdn.net/home/phong/Workspace/WhisperVQ/~/Workspace/WhisperVQ/.venv/lib/python3.11/site-packages/torch/nn/modules/module.py:1565)     result = None

File ~/Workspace/WhisperVQ/.venv/lib/python3.11/site-packages/whisperspeech/modules.py:129, in MultiHeadAttention.forward(self, qx, q_positions, kvx, kv_positions, causal, mask)
    [126](https://vscode-remote+attached-002dcontainer-002b7b22636f6e7461696e65724e616d65223a222f72657365617263685f6c6162227d-0040ssh-002dremote-002b7b22686f73744e616d65223a224a616e2d34784136303030227d.vscode-resource.vscode-cdn.net/home/phong/Workspace/WhisperVQ/~/Workspace/WhisperVQ/.venv/lib/python3.11/site-packages/whisperspeech/modules.py:126)     q,k,v = None,None,None
    [128](https://vscode-remote+attached-002dcontainer-002b7b22636f6e7461696e65724e616d65223a222f72657365617263685f6c6162227d-0040ssh-002dremote-002b7b22686f73744e616d65223a224a616e2d34784136303030227d.vscode-resource.vscode-cdn.net/home/phong/Workspace/WhisperVQ/~/Workspace/WhisperVQ/.venv/lib/python3.11/site-packages/whisperspeech/modules.py:128) if q is None: q = self.query(qx) * self.sqrt_qk_scale
--> [129](https://vscode-remote+attached-002dcontainer-002b7b22636f6e7461696e65724e616d65223a222f72657365617263685f6c6162227d-0040ssh-002dremote-002b7b22686f73744e616d65223a224a616e2d34784136303030227d.vscode-resource.vscode-cdn.net/home/phong/Workspace/WhisperVQ/~/Workspace/WhisperVQ/.venv/lib/python3.11/site-packages/whisperspeech/modules.py:129) q = self.split_heads(q, q_positions, rope = self.rotary, subsampling = self.query_subsampling)
    [131](https://vscode-remote+attached-002dcontainer-002b7b22636f6e7461696e65724e616d65223a222f72657365617263685f6c6162227d-0040ssh-002dremote-002b7b22686f73744e616d65223a224a616e2d34784136303030227d.vscode-resource.vscode-cdn.net/home/phong/Workspace/WhisperVQ/~/Workspace/WhisperVQ/.venv/lib/python3.11/site-packages/whisperspeech/modules.py:131) if kvx is not self.cached_kvx:
    [132](https://vscode-remote+attached-002dcontainer-002b7b22636f6e7461696e65724e616d65223a222f72657365617263685f6c6162227d-0040ssh-002dremote-002b7b22686f73744e616d65223a224a616e2d34784136303030227d.vscode-resource.vscode-cdn.net/home/phong/Workspace/WhisperVQ/~/Workspace/WhisperVQ/.venv/lib/python3.11/site-packages/whisperspeech/modules.py:132)     if k is None: k = self.key(kvx) * self.sqrt_qk_scale

File ~/Workspace/WhisperVQ/.venv/lib/python3.11/site-packages/whisperspeech/modules.py:108, in MultiHeadAttention.split_heads(self, x, x_positions, rope, subsampling)
    [106](https://vscode-remote+attached-002dcontainer-002b7b22636f6e7461696e65724e616d65223a222f72657365617263685f6c6162227d-0040ssh-002dremote-002b7b22686f73744e616d65223a224a616e2d34784136303030227d.vscode-resource.vscode-cdn.net/home/phong/Workspace/WhisperVQ/~/Workspace/WhisperVQ/.venv/lib/python3.11/site-packages/whisperspeech/modules.py:106) x = x.view(*x.shape[:2], self.n_head, -1)
    [107](https://vscode-remote+attached-002dcontainer-002b7b22636f6e7461696e65724e616d65223a222f72657365617263685f6c6162227d-0040ssh-002dremote-002b7b22686f73744e616d65223a224a616e2d34784136303030227d.vscode-resource.vscode-cdn.net/home/phong/Workspace/WhisperVQ/~/Workspace/WhisperVQ/.venv/lib/python3.11/site-packages/whisperspeech/modules.py:107) if rope:
--> [108](https://vscode-remote+attached-002dcontainer-002b7b22636f6e7461696e65724e616d65223a222f72657365617263685f6c6162227d-0040ssh-002dremote-002b7b22686f73744e616d65223a224a616e2d34784136303030227d.vscode-resource.vscode-cdn.net/home/phong/Workspace/WhisperVQ/~/Workspace/WhisperVQ/.venv/lib/python3.11/site-packages/whisperspeech/modules.py:108)     x = rope_rotate(x, x_positions * subsampling, *self.rotary(x))
    [109](https://vscode-remote+attached-002dcontainer-002b7b22636f6e7461696e65724e616d65223a222f72657365617263685f6c6162227d-0040ssh-002dremote-002b7b22686f73744e616d65223a224a616e2d34784136303030227d.vscode-resource.vscode-cdn.net/home/phong/Workspace/WhisperVQ/~/Workspace/WhisperVQ/.venv/lib/python3.11/site-packages/whisperspeech/modules.py:109) return x.permute(0, 2, 1, 3)

TypeError: unsupported operand type(s) for *: 'NoneType' and 'int'

Looking for help, thank you in advance!