code: def calc_variance(grad_dict, num_batches, param_names):
      
      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
docs: 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