MLPClassifierSpace on papalexi_2021 fails with TypeError: cannot pickle 'weakref.ReferenceType' object

[Lilly-May]

Report

When using the MLPClassifierSpace on the papalexi_2021 dataset, I get the following error during sanity checking:

TypeError: cannot pickle 'weakref.ReferenceType' object

The MLPClassifierSpace works fine for the dummy dataset during testing and for the Norman dataset, so I assume that the issue is not due to the implementation of the MLPClassifierSpace in general but may be caused by some characteristic of the dataset. Also, the error occurs during the fitting of the Trainer, and there are existing bug reports for similar issues with PyTorch Lightning, so I'm not sure how much can be done from our end to address it... I would appreciate it if others could attempt to reproduce the issue. If it persists for them, I'll open an issue on the PyTorch Lightning repository.

Here's the code to reproduce the error:

adata = pt.dt.papalexi_2021()["rna"]
dcs = pt.tl.MLPClassifierSpace()
dcs = dcs.load(adata, target_col="gene_target")
dcs.train()
cell_embeddings = dcs.get_embeddings()

Or with the new compute method from #565:

adata = pt.dt.papalexi_2021()["rna"]
dcs = pt.tl.MLPClassifierSpace()
cell_embeddings = dcs.compute(adata, target_col="gene_target")

Version information

pytorch-lightning: 2.0.9 torch: 2.0.1

[Zethson]

Could you please paste your lightning + pytorch versions?

[Zethson]

Currently checking with the latest versions....

pytorch-lightning-2.2.1

torch-2.2.1

[Zethson]

Can't reproduce this issue on my Linux system with the latest versions.

[Lilly-May]

Still persists for me even with the newest versions :/ But since this seems to be a bug specific to my machine, and it exists both before and after introducing the compute method in #565 and is therefore not directly linked to the PRs referenced above, I'd go ahead and mark the PRs ready for review. I'll then address this issue separately. What do you think @Zethson?

[Zethson]

Yes, please mark it ready for review.

We can also ask other people from #dream-team to narrow down what kind of Mac issue it is, but I doubt that we can do a lot about it. Maybe just report it to lightning in any case?

[Zethson]

I wonder whether you changed your Python versions inbetween or something?

I'll close this for now.