cannot pickle 'weakref.ReferenceType' object

[mousepixels]

Hi!

I am getting an error at every checkpoint and a warning about Jax being incompatible with os.fork()

It continues until the end and then gives an error about the checkpoint file not being saved (not shown).

Thanks!

Code:

cellbender remove-background \
--input temp/raw_adata.h5ad \
--output temp/out_adata.h5ad \
--expected-cells 9070 \
--total-droplets-included 50000  \
--cuda

Warning:

RuntimeWarning: os.fork() was called. os.fork() is incompatible with multithreaded code, and JAX is multithreaded, so this will likely lead to a deadlock.
  pid, fd = os.forkpty()

Error:

 File "python3.10/site-packages/cellbender/remove_background/checkpoint.py", line 115, in save_checkpoint
    torch.save(model_obj, filebase + '_model.torch')
  File "python3.10/site-packages/torch/serialization.py", line 629, in save
    _save(obj, opened_zipfile, pickle_module, pickle_protocol, _disable_byteorder_record)
  File "python3.10/site-packages/torch/serialization.py", line 841, in _save
    pickler.dump(obj)
TypeError: cannot pickle 'weakref.ReferenceType' object

System:

python 3.10 (tried in 3.11 too)
cellbender==0.3.0
jax==0.4.25
jaxlib==0.4.25
torch==2.2.1

os: Ubuntu 22

[mousepixels]

Missed the "install in own conda environment" bit. Works with fresh install in its own 3.7 environment. However, it would be nice for it to not conflict with scanpy/scvi-tools and be able to run this in a python script instead of from the command line only. But regardless, thanks for the great tool.