def normalize(data, trn_idx):
    """
    Normalize each dataset based on training index.
    """
    avg = data[trn_idx].mean(axis=0)
    std = data[trn_idx].std(axis=0)
    return (data - avg) / std