RuntimeError: istft(torch.cuda.FloatTensor[2, 1, 1539, 214, 2], ... : expected a tensor with 3 or 4 dimensions, but got 5

[madhavmk]

Thanks for the clear code implementation of the Complex Unet paper !

Unfortunately I'm facing a Tensor Dimension RunTime error when training the DCUnet on both my local setup and on Colaboratory. I have pasted the error stack below.

A Runtime Error occurs when running Cell 28 :
train_losses, test_losses = train(dcunet10, train_loader, test_loader, loss_fn, optimizer, scheduler, 3)

0% |          | 0/3 [00:00<?, ?it/s]/usr/local/lib/python3.6/dist-packages/torchaudio/functional.py:110: UserWarning: istft has been moved to PyTorch and will be removed from torchaudio, please use torch.istft instead.
'istft has been moved to PyTorch and will be removed from torchaudio, '
0% |          | 0/3 [00:01<?, ?it/s]
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-30-c06ae68fffa4> in <module>()
----> 1 train_losses, test_losses = train(dcunet10, train_loader, test_loader, loss_fn, optimizer, scheduler, 3)

5 frames
<ipython-input-25-6d7e3c9eda72> in train(net, train_loader, test_loader, loss_fn, optimizer, scheduler, epochs)
      9         if e == 0:
     10             with torch.no_grad():
---> 11                 test_loss = test_epoch(net, test_loader, loss_fn)
     12 
     13             test_losses.append(test_loss)

<ipython-input-23-310a2e7c0bd7> in test_epoch(net, test_loader, loss_fn)
      6         # get the output from the model
      7         noisy_x, clean_x = noisy_x.to(DEVICE), clean_x.to(DEVICE)
----> 8         pred_x = net(noisy_x)
      9 
     10         # calculate loss

/usr/local/lib/python3.6/dist-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
    720             result = self._slow_forward(*input, **kwargs)
    721         else:
--> 722             result = self.forward(*input, **kwargs)
    723         for hook in itertools.chain(
    724                 _global_forward_hooks.values(),

<ipython-input-20-4e83df3f0c6a> in forward(self, x, is_istft)
     53         output = u4 * x
     54         if is_istft:
---> 55             output = torchaudio.functional.istft(output, n_fft=self.n_fft, hop_length=self.hop_length, normalized=True)
     56 
     57         return output

/usr/local/lib/python3.6/dist-packages/torchaudio/functional.py in istft(stft_matrix, n_fft, hop_length, win_length, window, center, pad_mode, normalized, onesided, length)
    116     return torch.istft(
    117         input=stft_matrix, n_fft=n_fft, hop_length=hop_length, win_length=win_length, window=window,
--> 118         center=center, normalized=normalized, onesided=onesided, length=length)
    119 
    120 

/usr/local/lib/python3.6/dist-packages/torch/functional.py in istft(input, n_fft, hop_length, win_length, window, center, normalized, onesided, length)
    525 
    526     return _VF.istft(
--> 527         input, n_fft, hop_length, win_length, window, center, normalized, onesided, length)
    528 
    529 

RuntimeError: istft(torch.cuda.FloatTensor[2, 1, 1539, 214, 2], n_fft=3076, hop_length=772, win_length=3076, window=None, center=1, normalized=1, onesided=1, length=None): expected a tensor with 3 or 4 dimensions, but got 5

Any help on this would be appreciated. Thanks :)

[pheepa]

Hi, thank you for issue. Try to squeeze(1).

[madhavmk]

Yes. using 'torch.squeeze(var, 1)' to return tensors with single dimensional quantities removed seems to be working.
I will create a pull request with this fix once i verify that there are no errors.