search

havakv / pycox

Survival analysis with PyTorch
BSD 2-Clause "Simplified" License
822 stars 191 forks source link

Is there a way to save a trained model? #29

Closed bckwon closed 4 years ago

bckwon commented 4 years ago

Hi,

I was able to train a DeepSurv model following the tutorial, but I was wondering if I can save the trained model. I tried to pickle it, but it generates the following error:

AttributeError: Can't pickle local object '_ActionOnBestMetric.on_fit_start.<locals>.<lambda>'

I also attempted to use state_dict, but CoxPH doesn't have such an attribute.

In the end, I would like to create a checkpoint so that I can save and load models. Could you please let me know if there's a method? Thank you.

bckwon commented 4 years ago

I found that there are methods called save_net and load_net.

havakv commented 4 years ago

I'm happy you figured it out! There aren't really any docs now, which would make things like that much simpler.