Error saving / loading when initialize the model with kappa/ rho etc.

[ruochiz]

Hi,

Thank you for creating this amazing tool. I noticed a potential bug when saving/loading models.

File spectra.py Line: near 243 & 251. The dictionary self.rho / self.kappa should be transformed in to nn.ParameterDict after initialization to register the parameters to the state_dict. Otherwise, these parameters won't be saved, and hence raise an error when trying to load trained model from disk.

The same goes for cases. when use_cell_types = False, kappa and rho should be registered to be saved properly.

[wallet-maker]

Hi,

thank you for your feedback and your suggestions they are much appreciated. Yes, rho and kappa parameters currently have to be initialized identically when you load the model from disk. Do you happen to have the error message when you tried loading the model?

Thanks, Thomas

[ruochiz]

Error info is like this,

and it can be fixed by

The error happens when I run the code following example notebook and later try to reload the trained model.

Thanks.

[wallet-maker]

That's super helpful thank you, we're working to fix that!