diff --git a/CHANGELOGS.rst b/CHANGELOGS.rst
index ed4c1242..1cb81bbe 100644
--- a/CHANGELOGS.rst
+++ b/CHANGELOGS.rst
@@ -1,6 +1,11 @@
 Change Logs
 ===========
 
+0.9.4
++++++
+
+* :pr:`429`: fix issue related to PyTreeSpec
+
 0.9.3
 +++++
 
diff --git a/_doc/index.rst b/_doc/index.rst
index 4347b731..405d36f7 100644
--- a/_doc/index.rst
+++ b/_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:
diff --git a/onnx_diagnostic/__init__.py b/onnx_diagnostic/__init__.py
index b05d374e..2ab745b4 100644
--- a/onnx_diagnostic/__init__.py
+++ b/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é"
diff --git a/onnx_diagnostic/investigate/input_observer.py b/onnx_diagnostic/investigate/input_observer.py
index 94796d32..c7fd7eb8 100644
--- a/onnx_diagnostic/investigate/input_observer.py
+++ b/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
diff --git a/pyproject.toml b/pyproject.toml
index bf543d6b..99d063cd 100644
--- a/pyproject.toml
+++ b/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 = [