gimVI - multi-GPU support? #3057
Description
I'm currently running gimVI on 14K genes to impute spatial expression from scRNA-seq data. My model trains, but then I can't get the imputed expression using model.get_imputed_values(). My 1 GPU has ran out of memory.
imputed_spatial_rates = model.get_imputed_values(normalized= False)[1]OutOfMemoryError: CUDA out of memory. Tried to allocate 20.00 MiB. GPU 0 has a total capacity of 44.53 GiB of which 448.00 KiB is free. Including non-PyTorch memory, this process has 44.52 GiB memory in use. Of the allocated memory 36.12 GiB is allocated by PyTorch, and 7.91 GiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation. See documentation for Memory Management (https://pytorch.org/docs/stable/notes/cuda.html#environment-variables)```
#### Versions:
<!-- Output of scvi.__version__ -->
> VERSION
1.2.0
<!-- Relevant screenshots -->
Question: Can I run on multi-GPUs to get the learned expression for GIMVI to overcome this out-of-memory problem?
Thank you.!
[Image](https://github.com/user-attachments/assets/e6ae4441-56b2-4530-956f-d9710f60bb3b)