import torch


def normalize_data(batch_data, batch_std=None, batch_mean=None):
    """
    Normalize data using Standard Normalizer.

    Usage:
    normalized_data, batch_std, batch_mean = normalize_data(batch_data)

    Args:
        batch_data -> data to normalize
        batch_std -> Standard Deviation (If None, calculated automatically)
        batch_mean -> Mean (Average) (If None, calculated automatically)
    """
    if batch_std is None:
        batch_std = torch.std(batch_data, 0)
    if batch_mean is None:
        batch_mean = torch.mean(batch_data, 0)
    batch_data = (batch_data - batch_mean) / batch_std
    batch_data = torch.nan_to_num(batch_data)

    return batch_data, batch_std, batch_mean