import torch


def batch_mse_torch(estimation, origin):
    """
    batch-wise mse caculation for multiple audio files.
    estimation: (batch, nsource, frames, freq_bins)
    origin: (batch, nsource, frames, freq_bins)
    nsource = 2
    """
    mse1 = torch.sqrt(torch.pow(estimation - origin, 2).mean([3])).mean([1,2])
    mse2 = torch.sqrt(torch.pow(estimation - origin.flip([1]), 2).mean([3])).mean([1,2])
    return torch.stack((mse1, mse2),1).min(1)[0]