Commit 8728290d authored by Noric Couderc's avatar Noric Couderc
Browse files

Fixed saving the names of features 

The feature names were not actually saved originally.
parent 104e7937

train_model.py

+5 −3
Original line number Diff line number Diff line
@@ -368,7 +368,9 @@ if __name__ == "__main__": 


    classifier.fit(X_train, y_train, w_train)

    # We just add it to the class

    classifier.feature_names = poly_transformer.get_feature_names()



    classifier.feature_names = poly_transformer \

        .get_feature_names(features_extended.columns.values)



    print("Accuracy: {0}".format(classifier.score(X_test, y_test, w_test)))

    print("Classifier trained")