resolve document_type in DocumentStatistic.__init__

ArneBinder · Sep 14, 2023 · aa089f1 · aa089f1
1 parent cdbd5e6
commit aa089f1
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/src/pytorch_ie/core/statistic.py b/src/pytorch_ie/core/statistic.py
@@ -5,7 +5,7 @@
 
 from pytorch_ie.core.document import Document
 from pytorch_ie.core.metric import DocumentMetric
-from pytorch_ie.utils.hydra import InstantiationException, resolve_target
+from pytorch_ie.utils.hydra import InstantiationException, resolve_document_type, resolve_target
 
 logger = logging.getLogger(__name__)
 
@@ -152,7 +152,7 @@ def __init__(
         show_as_markdown: bool = False,
         aggregation_functions: Optional[List[str]] = None,
         title: Optional[str] = None,
-        document_type: Optional[Type[Document]] = None,
+        document_type: Optional[Union[Type[Document], str]] = None,
     ) -> None:
         super().__init__()
         self.aggregation_functions = {
@@ -162,7 +162,7 @@ def __init__(
         self.show_histogram = show_histogram
         self.show_as_markdown = show_as_markdown
         self.title = title or self.__class__.__name__
-        self._document_type = document_type
+        self._document_type = resolve_document_type(document_type)
 
     @property
     def document_type(self) -> Optional[Type[Document]]: