Fix metric_factory's type

neuro-galaxy · Nov 28, 2024 · 69230b9 · 69230b9
1 parent 73f7586
commit 69230b9
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/torch_brain/utils/stitcher.py b/torch_brain/utils/stitcher.py
@@ -115,7 +115,7 @@ def __init__(
         self,
         session_ids: Iterable[str],
         modality_spec: Optional[ModalitySpec] = None,
-        metric_factory: Optional[Callable[[int], ModalitySpec]] = None,
+        metric_factory: Optional[Callable[..., torchmetrics.Metric]] = None,
         quiet=False,
     ):
         r"""