BaseEncoderMaskerDecoder instance cannot be converted to TorchScript?

[urig]

🐛 Bug

When trying to convert a pre-trained instance of BaseEncoderMaskerDecoder to TorchScript a RuntimeError is thrown, seemingly because of incompatible code inside asteroid's Conv1DBlocks class.

To Reproduce

Run this bit of code where asteroid v0.6.0 is installed:

    from asteroid.models.base_models import BaseEncoderMaskerDecoder

    model = BaseEncoderMaskerDecoder.from_pretrained("JorisCos/ConvTasNet_Libri3Mix_sepnoisy_16k")
    model_scripted = torch.jit.script(model)

Expected behavior

model_scripted should contain the TorchScript equivalent of the model.

Actual behavior

A RuntimeError is raised from the last line of code:

Previous return statement returned a value of type Tensor but this return statement returns a value of type Tuple[Tensor, Tensor]:
  File "C:\dev\github\SilentiumIsrael\ResearchML\.venv\lib\site-packages\asteroid\masknn\convolutional.py", line 97
            return res_out
        skip_out = self.skip_conv(shared_out)
        return res_out, skip_out
        ~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE

Examining https://github.com/asteroid-team/asteroid/blob/v0.6.0/asteroid/masknn/convolutional.py#L90 indeed shows two return statements each returning a different type.

Environment

Package versions

Asteroid            0.6.0
PyTorch             1.9.0+cpu
PyTorch-Lightning   1.7.7

Additional info

Reproduced on Windows with Python 3.9.6.

[mpariente]

Thanks for opening the issue.

Have you looked further into this ? Digged the unit tests ? versions of Pytorch ? etc ?

[urig]

Hi @mpariente.
Thanks for responding.
The version of PyTorch is 1.9.1.
I'll try and clear some time to go over the tests and maybe create a failing test for this.
Cheers, Uri