🐛 Bug

There seems to be a significant performance drop after migrating the istft from torchaudio to pytorch.
migration discussed here #34827
I've also noticed the performance gap increases with greater n_fft.

To Reproduce

pip install torch==1.6.0

import torch
import time
n_fft = 2**15
x = torch.rand(size=[1, n_fft//2 + 1, 10, 2])
start = time.time()
y = torch.functional.istft(x,n_fft)
end = time.time()
print(end - start)

pip install torchaudio==0.5.1

import torch
import torchaudio
import time
n_fft = 2**15
x = torch.rand(size=[1, n_fft//2 + 1, 10, 2])
start = time.time()
y = torchaudio.functional.istft(x,n_fft)
end = time.time()
print(end - start)

Expected behavior

I would expect similar timings between:

torch.functional.istft in torch==1.6.0
and
torchaudio.functional.istft in torchaudio==0.5.1

Environment running torch.functional.istft

PyTorch version: 1.6.0
Is debug build: No
CUDA used to build PyTorch: 10.2

OS: Ubuntu 18.04.3 LTS
GCC version: (Ubuntu 7.5.0-3ubuntu1~18.04) 7.5.0
CMake version: version 3.12.0

Python version: 3.6
Is CUDA available: No
CUDA runtime version: 10.1.243
GPU models and configuration: Could not collect
Nvidia driver version: Could not collect
cuDNN version: /usr/lib/x86_64-linux-gnu/libcudnn.so.7.6.5

Versions of relevant libraries:
[pip3] numpy==1.18.5
[pip3] torch==1.6.0
[pip3] torchaudio==0.6.0
[pip3] torchsummary==1.5.1
[pip3] torchtext==0.3.1
[pip3] torchvision==0.6.1+cu101
[conda] Could not collect

Environment running torchaudio.functional.istft

PyTorch version: 1.5.1+cu101
Is debug build: No
CUDA used to build PyTorch: 10.1

OS: Ubuntu 18.04.3 LTS
GCC version: (Ubuntu 7.5.0-3ubuntu1~18.04) 7.5.0
CMake version: version 3.12.0

Python version: 3.6
Is CUDA available: No
CUDA runtime version: 10.1.243
GPU models and configuration: Could not collect
Nvidia driver version: Could not collect
cuDNN version: /usr/lib/x86_64-linux-gnu/libcudnn.so.7.6.5

Versions of relevant libraries:
[pip3] numpy==1.18.5
[pip3] torch==1.5.1+cu101
[pip3] torchaudio==0.5.1
[pip3] torchsummary==1.5.1
[pip3] torchtext==0.3.1
[pip3] torchvision==0.6.1+cu101
[conda] Could not collect

Additional context

Both environments are from google colab

cc @ezyang @gchanan @zou3519 @VitalyFedyunin @ngimel