def normalize(train, test):
    """
    Normalize the mean and variance
    
    Inputs:
        - train : training input features (N, D)
        - test  : test input features (M, D)
        
    Outputs:
        - train : normalized train set
        - test  : normalized test set (based on train meand and std)
    """
    mean, std = train.mean(dim=0), train.std(dim=0)
    train = (train - mean)/ std
    test = (test - mean)/std
    return train, test