Skip to content

Commit

Permalink
Changed Model Serialization to use CloudPickle, removed comment
Browse files Browse the repository at this point in the history
  • Loading branch information
vivekchettiar committed Dec 21, 2020
1 parent 943ae29 commit 2f3833d
Showing 1 changed file with 3 additions and 2 deletions.
5 changes: 3 additions & 2 deletions shap/models/_model.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
import pickle
import cloudpickle

class Model():
def __init__(self, model=None):
Expand All @@ -14,7 +15,7 @@ def __call__(self, *args):

def save(self, out_file, *args):
pickle.dump(type(self), out_file)
pickle.dump(self.model, out_file) # TODO change serialization methods based on model
cloudpickle.dump(self.model, out_file)

@classmethod
def load(cls, in_file):
Expand All @@ -23,4 +24,4 @@ def load(cls, in_file):
print("Warning: model was not found in saved file, please set model before using explainer.")
return None

return pickle.load(in_file)
return cloudpickle.load(in_file)

0 comments on commit 2f3833d

Please sign in to comment.