The former iSTFT code has a bug which will produce wrong results when batch_size > 1:

coffidx = th.where(coff > 1e-8)
outputs[coffidx] = outputs[coffidx]/(coff[coffidx])

When you use torch.where(condition), you get a tuple containing 3 tensors representing the indices on dimension 0, 1, 2, respectively. However, coffidx is of shape [1, 1, T] while outputs is of shape [B, 1, T]. Thus, you will only get 0 on the first dim of coff and the later calculation only happens on outputs[0, ...]. The rest part (output[1:, ...]) is never updated.

The fix is either using torch.repeat to make coff [B, 1, T] too, or using torch.where(condition, a, b) to get a combined coff based on the condition. This PR applies the latter solution.