Use joblib to save the model

train.py

@@ -5,14 +5,14 @@ stores the model in a pickle file.
 
 """
 
-import pickle
-
+import joblib
 from sklearn.datasets import make_classification
 from sklearn.linear_model import SGDClassifier
 from sklearn.model_selection import GridSearchCV
 
 
 SEED = 0
+FILENAME = 'sklearn_model.joblib'
 
 
 def get_data():
@@ -48,8 +48,7 @@ def train(model, X, y, hparams):
 
 
 def save_model(model, filename):
-    with open(filename, 'wb') as f:
-        pickle.dump(model, f)
+    joblib.dump(model, filename)
     print(f"Stored model in '{filename}'")
 
 
@@ -58,7 +57,7 @@ def main():
     model = get_model()
     hparams = get_hparams()
     model_trained = train(model, X, y, hparams=hparams)
-    save_model(model_trained, 'model.pickle')
+    save_model(model_trained, FILENAME)
 
 
 if __name__ == '__main__':