torch.jit.frontend.UnsupportedNodeError: DictComp aren't supported

I know it's not strictly related to Resemblyzer library, but I do hope that may be you or others can give me an help.

My objective is to use the following class which simply uses your excellent work:

import torch
from torch.utils.data import DataLoader
import torchaudio
from resemblyzer import preprocess_wav, VoiceEncoder
from demo_utils import *
from itertools import groupby
from pathlib import Path
from tqdm import tqdm
import numpy as np
import sys

np.set_printoptions(threshold=sys.maxsize)

class GetSpeechEmbedding(torch.nn.Module):
    def __init__(self):
        super(GetSpeechEmbedding, self).__init__()

    def forward(self, a, b, extension):
        encoder = VoiceEncoder()
        p = Path(a,b)
        wav_fpaths = list(p.glob("**/*." + extension))
        speaker_wavs = {speaker: list(map(preprocess_wav, wav_fpaths)) for speaker, wav_fpaths 
in
                    groupby(tqdm(wav_fpaths, "Preprocessing wavs", len(wav_fpaths), unit="wavs"), 
                            lambda wav_fpath: wav_fpath.parent.stem)}
        embeds_a = np.array([encoder.embed_utterance(wavs[0]) for wavs in 
speaker_wavs.values()])
        return embeds_a

With python it works fine:

(venv373) (base) marco@pc01:~/PyTorchMatters/Resemblyzer$ python3 
./getSpeechEmbedding.py 
Loaded the voice encoder model on cpu in 0.03 seconds.
Preprocessing wavs:     
 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 100/100 [00:01<00:00, 78.71wavs/s]
[[0.00000000e+00 4.73524793e-04 1.30652651e-01 0.00000000e+00
 8.16560537e-02 8.52451399e-02 0.00000000e+00 0.00000000e+00
  7.62926340e-02 1.80580884e-01 2.26451114e-01 1.55212656e-01
  1.29384667e-01 0.00000000e+00 1.10790301e-04 0.00000000e+00

Following the indications found here: https://pytorch.org/tutorials/advanced/cpp_export.html

I added to getSpeechEmbedding.py these lines:

my_module = GetSpeechEmbedding()
sm = torch.jit.script(my_module)
sm.save("annotated_get_speech_embedding.pt")

But when trying to serialize the module, I get this error: "torch.jit.frontend.UnsupportedNodeError: DictComp aren't supported"

(venv373) (base) marco@pc01:~/PyTorchMatters/Resemblyzer$ python3 
./getSpeechEmbedding.py 
Traceback (most recent call last):
  File "./getSpeechEmbedding.py", line 114, in <module>
    sm = torch.jit.script(my_module)
  File "/home/marco/anaconda3/lib/python3.7/site-packages/torch/jit/__init__.py", line 1255, in 
script
    return torch.jit._recursive.recursive_script(obj)
  File "/home/marco/anaconda3/lib/python3.7/site-packages/torch/jit/_recursive.py", line 534, in 
recursive_script
    return create_script_module(nn_module, infer_methods_to_compile(nn_module))
  File "/home/marco/anaconda3/lib/python3.7/site-packages/torch/jit/_recursive.py", line 493, in 
infer_methods_to_compile
    stubs.append(make_stub_from_method(nn_module, method))
  File "/home/marco/anaconda3/lib/python3.7/site-packages/torch/jit/_recursive.py", line 40, in 
make_stub_from_method
    return make_stub(func)
  File "/home/marco/anaconda3/lib/python3.7/site-packages/torch/jit/_recursive.py", line 33, in 
make_stub
    ast = torch.jit.get_jit_def(func, self_name="RecursiveScriptModule")
  File "/home/marco/anaconda3/lib/python3.7/site-packages/torch/jit/frontend.py", line 171, in 
get_jit_def
    return build_def(ctx, py_ast.body[0], type_line, self_name)
  File "/home/marco/anaconda3/lib/python3.7/site-packages/torch/jit/frontend.py", line 212, in 
build_def
    build_stmts(ctx, body))
  File "/home/marco/anaconda3/lib/python3.7/site-packages/torch/jit/frontend.py", line 127, in 
build_stmts
    stmts = [build_stmt(ctx, s) for s in stmts]
  File "/home/marco/anaconda3/lib/python3.7/site-packages/torch/jit/frontend.py", line 127, in 
<listcomp>
    stmts = [build_stmt(ctx, s) for s in stmts]
  File "/home/marco/anaconda3/lib/python3.7/site-packages/torch/jit/frontend.py", line 187, in 
__call__
    return method(ctx, node)
  File "/home/marco/anaconda3/lib/python3.7/site-packages/torch/jit/frontend.py", line 289, in 
build_Assign
    rhs = build_expr(ctx, stmt.value)
  File "/home/marco/anaconda3/lib/python3.7/site-packages/torch/jit/frontend.py", line 186, in 
__call__
    raise UnsupportedNodeError(ctx, node)
torch.jit.frontend.UnsupportedNodeError: DictComp aren't supported:
  File "./getSpeechEmbedding.py", line 46
        # It normalizes the volume, trims long silences and resamples the wav to the correct 
sampling rate.

        speaker_wavs = {speaker: list(map(preprocess_wav, wav_fpaths)) for speaker, wav_fpaths
in
                       ~ <--- HERE
                    groupby(tqdm(wav_fpaths, "Preprocessing wavs", len(wav_fpaths), unit="wavs"), 
                        lambda wav_fpath: wav_fpath.parent.stem)}

I read here: https://pytorch.org/tutorials/advanced/cpp_export.html that "If you need to exclude some methods in your nn.Module because they use Python features that TorchScript doesn’t support yet, you could annotate those with @torch.jit.ignore" But I guess that excluding the dictionary comprehension from the serialization, will interfere with the method's functionality. What would you suggest me to do?

resemble-ai / Resemblyzer

torch.jit.frontend.UnsupportedNodeError: DictComp aren't supported #25