import torch


def calc_mse_loss(y_pred, y_true, grid_dims=(1,2)):
    """
    Calculates mean squared error 
    mse_loss = avg((y_pred)^2 - (y_true)^2)

    Args:
        y torch.tensor(n_samples, xgrid1, xgrid2, n_ydim)
        y_pred torch.tensor(n_samples, xgrid1, xgrid2, n_ydim)
        grid_dims tuple(): Indices of grid dimensions. Infinity error is taken across these dimensions
    Returns:
        mse_loss torch.tensor(n_samples, n_ydim)
    """
    n_px = 1
    for grid_dim in grid_dims:
        n_px = n_px * y_true.shape[grid_dim]
    mse_loss = 1./float(n_px)*(torch.sqrt(torch.sum(torch.pow(y_true - y_pred,2), dim=grid_dims)))# (n_samples,n_ch)
    return mse_loss