streaming_demo.py crashes with shape mismatch

[cobalamin]

Hi there! Thank you for providing code and checkpoints for your work, especially the streaming implementation :)

Unfortunately at the moment the streaming_demo.py is not working on my machine. It runs into the following error immediately when running the recommended command PYTHONPATH=./ python src/examples/streaming_demo.py:

An unexpected error occurred: The size of tensor a (295) must match the size of tensor b (63) at non-singleton dimension 1
Traceback (most recent call last):
  File "/export/home/.../CleanUMamba/src/examples/streaming_demo.py", line 135, in <module>
    model_output = model.feed(tensor_chunk)
                   ^^^^^^^^^^^^^^^^^^^^^^^^
  File "/export/home/.../CleanUMambaEnv/lib/python3.11/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/export/home/.../src/network/CleanUMamba.py", line 403, in feed
    out = self._denoise_frame(frame)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/export/home/.../src/network/CleanUMamba.py", line 477, in _denoise_frame
    x += skip_i[..., :x.shape[-1]]
RuntimeError: The size of tensor a (295) must match the size of tensor b (63) at non-singleton dimension 1

I did not modify the script or any parameters, keeping everything default. Might there a a configuration mismatch between the default checkpoint and the chunksize values etc. configured in the script, causing this shape mismatch? I don't get any error from PyAudio, so I'm assuming that it is accepting the sampling rate of 16000 Hz for my local soundcard.

[cobalamin]

I think your error is here in _denoise_frame of the network itself:

for i, upsampling_block in enumerate(self.decoder):
    skip_i = skip_connections[i]
    x += skip_i[..., :x.shape[-1]]

since the depth order is reversed on the upward path, but you collected the skip_connections from left to right on the downward path, it should read

skip_i = skip_connections[-(i+1)]

instead. That seems to fix the issue for me, at least mechanically.