Calculates the variance of the gradients per epoch for each parameter w.r.t number of batches

    Parameters
    ----------
    grad_dict: dict
        dictionary that maps parameter name to gradients in the mod executor group
    num_batches: int
        number of batches
    param_names: str
        parameter name in the module

    Returns
    ----------
    grad_dict: dict
        dictionary with new keys mapping to gradients variance

def calc_variance(grad_dict, num_batches, param_names):
    """Calculates the variance of the gradients per epoch for each parameter w.r.t number of batches

    Parameters
    ----------
    grad_dict: dict
        dictionary that maps parameter name to gradients in the mod executor group
    num_batches: int
        number of batches
    param_names: str
        parameter name in the module

    Returns
    ----------
    grad_dict: dict
        dictionary with new keys mapping to gradients variance

    """
    for i in range(len(param_names)):
        diff_sqr = mx.ndarray.square(mx.nd.subtract(grad_dict[param_names[i]],
                                                    grad_dict[str.format(param_names[i]+"_expectation")]))
        grad_dict[str.format(param_names[i] + "_variance")] = mx.ndarray.sum(diff_sqr, axis=0) / num_batches