fix torch import

scverse · Oct 22, 2024 · a0c7083 · a0c7083
1 parent 17cfd6b
commit a0c7083
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 3 deletions.
diff --git a/docs/conf.py b/docs/conf.py
@@ -73,9 +73,10 @@
 # Generate the API documentation when building
 autosummary_generate = True
 autodoc_member_order = "bysource"
+autodoc_mock_imports = ["torch"]
+# autodoc_default_flags = ['members']
 issues_github_path = "scverse/anndata"
 rtd_links_prefix = PurePosixPath("src")
-# autodoc_default_flags = ['members']
 napoleon_google_docstring = False
 napoleon_numpy_docstring = True
 napoleon_include_init_with_doc = False

diff --git a/src/anndata/experimental/pytorch/_annloader.py b/src/anndata/experimental/pytorch/_annloader.py
@@ -22,11 +22,12 @@
 
 if TYPE_CHECKING:
     from collections.abc import Callable, Generator, Sequence
-    from typing import TypeAlias
+    from typing import TypeAlias, Union
 
     from scipy.sparse import spmatrix
 
-    Array: TypeAlias = torch.Tensor | np.ndarray | spmatrix
+    # need to use Union because of autodoc_mock_imports
+    Array: TypeAlias = Union[torch.Tensor, np.ndarray, spmatrix]  # noqa: UP007
 
 
 # Custom sampler to get proper batches instead of joined separate indices