Pickle error when try to save a trained mode

I use pycox repository on Github for my project that is about survival analysis. This repository is based on Pytorch and torch tuples. After training a model we can save the model with


But when I use this code I face Pickle error that is not usual for other people. This error:

File ~/anaconda3/envs/myenv/lib/python3.9/site-packages/torchtuples/base.py:681 in save_net
    return torch.save(self.net, path, **kwargs)

  File ~/anaconda3/envs/myenv/lib/python3.9/site-packages/torch/serialization.py:380 in save
    _save(obj, opened_zipfile, pickle_module, pickle_protocol)

  File ~/anaconda3/envs/myenv/lib/python3.9/site-packages/torch/serialization.py:589 in _save

AttributeError: Can't pickle local object

Can anyone help me what could be the reason for this error?

I don’t know how save_net is implemented, but I would recommend to store the state_dict of the model instead of trying to pickle the model directly as it can fail in multiple ways during the loading (e.g. if you’ve changed the file structure or any function/class definition).

Because I use a repository and I do not save the model myself it is hard for me to understand how can I save it with state_dict. This is the save_net part of code in repository:

 def save_net(self, path, **kwargs):
        """Save self.net and baseline hazards to file.
            path {str} -- Path to file.
            **kwargs are passed to torch.save
        path, extension = os.path.splitext(path)
        if extension == "":
            extension = '.pt'
        super().save_net(path+extension, **kwargs)
        if hasattr(self, 'baseline_hazards_'):