def test_multi_scale_spectral_shape(batch_size):
# Test in with reduced number of STFT scales.
filt_list = [512, 256, 32]
# Fake targets and estimates
targets = torch.randn(batch_size, 8000)
est_targets = torch.randn(batch_size, 8000)
# Create PITLossWrapper in 'pw_pt' mode
loss_func = SingleSrcMultiScaleSpectral(windows_size=filt_list,
n_filters=filt_list,
hops_size=filt_list)
# Compute the loss
loss = loss_func(targets, est_targets)
assert loss.shape[0] == batch_size
def test_multi_scale_spectral_PIT(n_src):
# Test in with reduced number of STFT scales.
filt_list = [512, 256, 32]
# Fake targets and estimates
targets = torch.randn(2, n_src, 8000)
est_targets = torch.randn(2, n_src, 8000)
# Create PITLossWrapper in 'pw_pt' mode
pt_loss = SingleSrcMultiScaleSpectral(windows_size=filt_list,
n_filters=filt_list,
hops_size=filt_list)
loss_func = PITLossWrapper(pt_loss, pit_from='pw_pt')
# Compute the loss
loss = loss_func(targets, est_targets)