wrong import

CHANGELOGS.rst

@@ -1,6 +1,11 @@
 Change Logs
 ===========
 
+0.9.4
++++++
+
+* :pr:`429`: fix issue related to PyTreeSpec
+
 0.9.3
 +++++
 

_doc/index.rst

@@ -241,8 +241,8 @@ The function replaces dynamic dimensions defined as strings by
 Older versions
 ==============
 
+* `0.9.4 <../v0.9.4/index.html>`_
 * `0.9.3 <../v0.9.3/index.html>`_
-* `0.9.2 <../v0.9.2/index.html>`_
 * `0.8.11 <../v0.8.11/index.html>`_
 
 The documentation was updated on:

onnx_diagnostic/__init__.py

@@ -3,5 +3,5 @@
 Functions, classes to dig into a model when this one is right, slow, wrong...
 """
 
-__version__ = "0.9.3"
+__version__ = "0.9.4"
 __author__ = "Xavier Dupré"

onnx_diagnostic/investigate/input_observer.py

@@ -142,8 +142,8 @@ def __init__(
                 self.flat_list, self.spec
             )
 
-        self.aligned_spec: torch.utils._pytree.PyTreeSpec | None = None
-        self.aligned_flat_list: list[torch.Tensor | None] | None = None
+        self.aligned_spec = None
+        self.aligned_flat_list = None
 
     def remove_inputs(self, input_names: Sequence[str | int]):
         """Removes inputs."""
@@ -230,7 +230,7 @@ def n_tensors_for_args_kwargs(self) -> dict[int | str, int]:
     def _set_aligned_flat_list(
         self,
         aligned_flat_list: list[torch.Tensor | None],
-        aligned_spec: torch.utils._pytree.PyTreeSpec,
+        aligned_spec: "PyTreeSpec",  # noqa: F821
     ):
         self.aligned_flat_list = aligned_flat_list
         self.aligned_spec = aligned_spec
@@ -339,7 +339,7 @@ def __init__(
         self.default_values = default_values
         self.value_if_missing = value_if_missing
         self.inputs: list[InputCandidate] = []
-        self.outputs_specs: list[torch.utils._pytree.PyTreeSpec] = []
+        self.outputs_specs: list["PyTreeSpec"] = []  # noqa: F821, UP037
         self.flat_outputs: list[list[torch.Tensor | None]] = []
         self.latencies: list[float] = []
         self.args_name_and_position = args_name_and_position

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "onnx-diagnostic"
-version = "0.9.3"
+version = "0.9.4"
 description = "Tools to help converting pytorch models into ONNX."
 readme = "README.rst"
 authors = [