fix cuda test file

scverse · Dec 3, 2024 · 6f58e43 · 6f58e43
1 parent 5245163
commit 6f58e43
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/.github/workflows/test_linux_cuda.yml b/.github/workflows/test_linux_cuda.yml
@@ -59,7 +59,7 @@ jobs:
         run: |
           python -m pip install --upgrade pip wheel uv
           python -m uv pip install --system "scvi-tools[tests] @ ."
-          python -m pip install jax[cuda]
+          python -m pip install jax[cuda12]
           python -m pip install nvidia-nccl-cu12
 
       - name: Run pytest

diff --git a/tests/model/test_scvi.py b/tests/model/test_scvi.py
@@ -1300,7 +1300,7 @@ def test_scvi_num_workers():
     model.get_normalized_expression(n_samples=2)
 
 
-def test_scvi_train_ddp(save_path: str):
+def test_scvi_train_ddp(save_path: str = "."):
     training_code = """
 import torch
 import scvi
@@ -1312,7 +1312,7 @@ def test_scvi_train_ddp(save_path: str):
 model = SCVI(adata)
 
 model.train(
-    max_epochs=100,
+    max_epochs=1,
     check_val_every_n_epoch=1,
     accelerator="gpu",
     devices=-1,