fix _internal_output_axes

Author: FynnBe

bioimageio/core/model_adapters/_tensorflow_model_adapter.py

@@ -4,11 +4,12 @@
 
 import numpy as np
 
-from bioimageio.core.tensor import Tensor
 from bioimageio.spec.common import FileSource
 from bioimageio.spec.model import v0_4, v0_5
 from bioimageio.spec.utils import download
 
+from ..digest_spec import get_axes_infos
+from ..tensor import Tensor
 from ._model_adapter import ModelAdapter
 
 try:
@@ -73,11 +74,7 @@ def __init__(
         weight_file = self.require_unzipped(weights.source)
         self._network = self._get_network(weight_file)
         self._internal_output_axes = [
-            (
-                tuple(out.axes)
-                if isinstance(out.axes, str)
-                else tuple(a.id for a in out.axes)
-            )
+            tuple(a.id for a in get_axes_infos(out))
             for out in model_description.outputs
         ]
 

bioimageio/core/model_adapters/_torchscript_model_adapter.py

@@ -8,7 +8,7 @@
 from bioimageio.spec.model import v0_4, v0_5
 from bioimageio.spec.utils import download
 
-from ..axis import AxisId
+from ..digest_spec import get_axes_infos
 from ..tensor import Tensor
 from ._model_adapter import ModelAdapter
 
@@ -45,12 +45,10 @@ def __init__(
                 "Multiple devices for single torchscript model not yet implemented"
             )
 
-        self._model = torch.jit.load(  # pyright: ignore[reportPrivateImportUsage]
-            weight_path
-        )
+        self._model = torch.jit.load(weight_path)
         self._model.to(self.devices[0])
         self._internal_output_axes = [
-            tuple(AxisId(a) if isinstance(a, str) else a.id for a in out.axes)
+            tuple(a.id for a in get_axes_infos(out))
             for out in model_description.outputs
         ]