RuntimeError: "indices should be either on cpu or on the same device as the indexed tensor" during scnym_api training with GPU

[viraj-rapolu]

I am encountering a RuntimeError related to mismatched tensor devices during training when using scnym_api with GPU. The error occurs when indexing tensors during the training process, and it seems like the tensors are not being moved to the same device (CPU vs GPU) internally.

I have ensured that my data and configuration are set correctly, and the issue persists even when explicitly specifying the device in the config. This appears to be related to device management within the scnym library itself.

An example of my code that produces the issue:

`
import scnym
import scanpy as sc
import torch

adata = sc.read_h5ad('/path/to/dataset.h5ad')

sc.pp.normalize_total(adata)
sc.pp.log1p(adata)
sc.pp.highly_variable_genes(adata, n_top_genes=3000)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

scnym.scnym_api(
adata=adata, task='train', groupby='annotations',
domain_groupby='domain', out_path='/path/to/scnym_outs',
config={'device': str(device), 'other_config': 'no_new_identity'})
)
`

Logs/Traceback:

Traceback (most recent call last): File "run_scnym.py", line 102, in <module> config={'device': str(device), 'other_config': 'no_new_identity'}) File "/path/.local/lib/python3.7/site-packages/scnym/api.py", line 341, in scnym_api config=config, File "/path/.local/lib/python3.7/site-packages/scnym/api.py", line 538, in scnym_train **config['model_kwargs'], File "/path/.local/lib/python3.7/site-packages/scnym/main.py", line 519, in fit_model T.train() File "/path/.local/lib/python3.7/site-packages/scnym/trainer.py", line 452, in train self.train_epoch() File "/path/.local/lib/python3.7/site-packages/scnym/trainer.py", line 585, in train_epoch labeled_sample=data, File "/path/.local/lib/python3.7/site-packages/scnym/trainer.py", line 1491, in __call__ conf_unlabeled_sample[k] = unlabeled_sample[k][pseudolabel_confidence] RuntimeError: indices should be either on cpu or on the same device as the indexed tensor (cpu)

Environment information:

Name Version Build Channel
_libgcc_mutex 0.1 conda_forge conda-forge
_openmp_mutex 4.5 1_gnu conda-forge
absl-py 0.14.0 pypi_0 pypi
anndata 0.7.6 pypi_0 pypi
ca-certificates 2021.5.30 ha878542_0 conda-forge
cachetools 4.2.4 pypi_0 pypi
certifi 2021.5.30 pypi_0 pypi
charset-normalizer 2.0.6 pypi_0 pypi
configargparse 1.5.2 pypi_0 pypi
cycler 0.10.0 pypi_0 pypi
decorator 5.1.0 pypi_0 pypi
filelock 3.15.4 pypi_0 pypi
fsspec 2024.6.1 pypi_0 pypi
google-auth 1.35.0 pypi_0 pypi
google-auth-oauthlib 0.4.6 pypi_0 pypi
grpcio 1.41.0 pypi_0 pypi
h5py 3.4.0 pypi_0 pypi
idna 3.2 pypi_0 pypi
jinja2 3.1.4 pypi_0 pypi
joblib 1.0.1 pypi_0 pypi
kiwisolver 1.3.2 pypi_0 pypi
ld_impl_linux-64 2.36.1 hea4e1c9_2 conda-forge
leidenalg 0.8.7 pypi_0 pypi
libffi 3.4.2 h9c3ff4c_4 conda-forge
libgcc-ng 11.2.0 h1d223b6_9 conda-forge
libgomp 11.2.0 h1d223b6_9 conda-forge
libstdcxx-ng 11.2.0 he4da1e4_9 conda-forge
libzlib 1.2.11 h36c2ea0_1012 conda-forge
llvmlite 0.37.0 pypi_0 pypi
markdown 3.3.4 pypi_0 pypi
markupsafe 2.1.5 pypi_0 pypi
matplotlib 3.4.3 pypi_0 pypi
mock 4.0.3 pypi_0 pypi
mpmath 1.3.0 pypi_0 pypi
natsort 7.1.1 pypi_0 pypi
ncurses 6.2 h58526e2_4 conda-forge
networkx 2.6.3 pypi_0 pypi
numba 0.54.0 pypi_0 pypi
numexpr 2.7.3 pypi_0 pypi
numpy 1.20.3 pypi_0 pypi
nvidia-cublas-cu12 12.1.3.1 pypi_0 pypi
nvidia-cuda-cupti-cu12 12.1.105 pypi_0 pypi
nvidia-cuda-nvrtc-cu12 12.1.105 pypi_0 pypi
nvidia-cuda-runtime-cu12 12.1.105 pypi_0 pypi
nvidia-cudnn-cu12 9.1.0.70 pypi_0 pypi
nvidia-cufft-cu12 11.0.2.54 pypi_0 pypi
nvidia-curand-cu12 10.3.2.106 pypi_0 pypi
nvidia-cusolver-cu12 11.4.5.107 pypi_0 pypi
nvidia-cusparse-cu12 12.1.0.106 pypi_0 pypi
nvidia-nccl-cu12 2.20.5 pypi_0 pypi
nvidia-nvjitlink-cu12 12.6.20 pypi_0 pypi
nvidia-nvtx-cu12 12.1.105 pypi_0 pypi
oauthlib 3.1.1 pypi_0 pypi
openssl 3.0.0 h7f98852_1 conda-forge
packaging 21.0 pypi_0 pypi
pandas 1.3.3 pypi_0 pypi
patsy 0.5.2 pypi_0 pypi
pillow 10.4.0 pypi_0 pypi
pip 21.2.4 pyhd8ed1ab_0 conda-forge
protobuf 3.18.0 pypi_0 pypi
pyasn1 0.4.8 pypi_0 pypi
pyasn1-modules 0.2.8 pypi_0 pypi
pynndescent 0.5.4 pypi_0 pypi
pyparsing 2.4.7 pypi_0 pypi
python 3.9.7 hf930737_3_cpython conda-forge
python-dateutil 2.8.2 pypi_0 pypi
python-igraph 0.9.6 pypi_0 pypi
python_abi 3.9 2_cp39 conda-forge
pytz 2021.1 pypi_0 pypi
pyyaml 5.4.1 pypi_0 pypi
readline 8.1 h46c0cb4_0 conda-forge
requests 2.26.0 pypi_0 pypi
requests-oauthlib 1.3.0 pypi_0 pypi
rsa 4.7.2 pypi_0 pypi
scanpy 1.8.1 pypi_0 pypi
scikit-learn 1.0 pypi_0 pypi
scipy 1.7.1 pypi_0 pypi
scnym 0.1.11 pypi_0 pypi
seaborn 0.11.2 pypi_0 pypi
setuptools 58.0.4 py39hf3d152e_2 conda-forge
sinfo 0.3.4 pypi_0 pypi
six 1.16.0 pypi_0 pypi
sqlite 3.36.0 h9cd32fc_2 conda-forge
statsmodels 0.13.0rc0 pypi_0 pypi
stdlib-list 0.8.0 pypi_0 pypi
sympy 1.13.1 pypi_0 pypi
tables 3.6.1 pypi_0 pypi
tensorboard 2.6.0 pypi_0 pypi
tensorboard-data-server 0.6.1 pypi_0 pypi
tensorboard-plugin-wit 1.8.0 pypi_0 pypi
texttable 1.6.4 pypi_0 pypi
threadpoolctl 3.0.0 pypi_0 pypi
tk 8.6.11 h27826a3_1 conda-forge
torch 2.4.0 pypi_0 pypi
torchvision 0.19.0 pypi_0 pypi
tqdm 4.62.3 pypi_0 pypi
triton 3.0.0 pypi_0 pypi
typing-extensions 4.12.2 pypi_0 pypi
tzdata 2021a he74cb21_1 conda-forge
umap-learn 0.5.1 pypi_0 pypi
urllib3 1.26.7 pypi_0 pypi
werkzeug 2.0.1 pypi_0 pypi
wheel 0.37.0 pyhd8ed1ab_1 conda-forge
xlrd 1.2.0 pypi_0 pypi
xz 5.2.5 h516909a_1 conda-forge
zlib 1.2.11 h36c2ea0_1012 conda-forge

[jacobkimmel]

Hey Viraj -- sorry you ran into this bug! We haven't seen this before to my knowledge.

A couple thoughts:

1. scnym doesn't actually support a device kwarg in the config. It current assumes that you want to use a CUDA device if you have one available and moves everything over there. The device kwarg won't throw an error, but it won't be used.
2. Likewise, the config kwarg on the API call only accepts strings for pre-registered configs (see scnym.api.CONFIGS.keys()) or a full dictionary of configuration options that matches the format of those pre-registered options. Apologies, we never fully built out the config system for custom research, so everything outside the pre-registered configs is currently unsupported.
3. If you try running the following code, what results do you get?

import scnym
import scanpy as sc
import torch

adata = scanpy.datasets.pbmc3k()
# create some random class labels, the model should still hit high performance by overfitting
adata.obs["annotations"] = np.random.randint(0, 3, size=adata.shape[0])

# NOTE: scNym requires data to be log1p(CountsPerMillion), so we set `target_sum = 1e6` rather than the default `1e4`
sc.pp.normalize_total(adata, target_sum=int(1e6)) 
sc.pp.log1p(adata)
# cut the highly variable genes call -- scNym selects genes internally
# without subsetting the AnnData like 
# `adata = adata[:, adata.var["highly_variable"]], this doesn't have the intended effect anyway
# sc.pp.highly_variable_genes(adata, n_top_genes=3000)

scnym.scnym_api(
  adata=adata, 
  task='train', 
  groupby='annotations',
  domain_groupby='domain', 
  out_path='/path/to/scnym_outs',
  config="no_new_identity",
)

[viraj-rapolu]

Hi Jacob,

Thanks for looking into this so quickly! I modified the preprocessing steps and used the parameters you suggested for the scnym_api function, but I'm still encountering the same error. I'm having this issue even if I don't explicitly set the torch.device().