device="cuda" hardcoded in emb_extractor.py and perturber_utils.py, breaks CPU-only inference
When running EmbExtractor.extract_embs() on a CPU-only machine (no NVIDIA GPU), the following error occurs:
AssertionError: Torch not compiled with CUDA enabled
This is because device="cuda" is hardcoded in three places:
geneformer/emb_extractor.py, line 101: torch.tensor(minibatch["length"], device="cuda")geneformer/emb_extractor.py, line 111: input_data_minibatch.to("cuda")geneformer/perturber_utils.py, line 751: torch.tensor(attention_mask, device="cuda")
A simple fix that resolves this for CPU-only use: replace these three hardcoded "cuda" references with "cpu" (or better, a dynamic check like "cuda" if torch.cuda.is_available() else "cpu", which the codebase already uses correctly elsewhere, e.g. perturber_utils.py line 194).
Is this the intended behavior, or worth patching upstream? Happy to open a PR with this fix if useful.
Thank you for bringing this up. The default is expected GPU usage given it is much more efficient. The dynamic check is a great suggestion, and we would appreciate if you opened a PR with this fix!