def __init__(
        self,
        wrapped_type: type,
        name: Optional[str] = None,
    ):
        """
        Specify a model state as a wrapper of a ``TensorType``.
        For example, you can use the following code to create a
        state type input that wraps a fp16 tensor with shape ``(2, 3)``::

            ct.StateType(
                wrapped_type=ct.TensorType(
                    shape=(2, 3),
                    dtype=np.float16
                ),
                name="state",
            )

        Parameters
        ----------
        wrapped_type: coremltools.converters.mil.input_types.InputType
            - The type wrapped in the state. 
            - Must be ``TensorType``.
              Note that the ``name`` and ``default_value`` of the wrapped ``TensorType`` must not be provided.

        name: str
            The name of the state.
            It must match the key of ``named_buffers()`` in the source TorchScript model.
        """
        if not isinstance(wrapped_type, StateType.SUPPORTED_WRAPPER_TYPE):
            raise ValueError(
                f"StateType only supports {StateType.SUPPORTED_WRAPPER_TYPE}. Got {type(wrapped_type)}."
            )
        # name and default_value cannot be set
        if wrapped_type.name is not None:
            raise ValueError("name cannot be set in the state wrapped_type.")
        if wrapped_type.default_value is not None:
            raise ValueError("default_value cannot be set in the state wrapped_type.")

        super(StateType, self).__init__(name, wrapped_type.shape, wrapped_type.dtype)
        self.wrapped_type = wrapped_type

        Specify a model state as a wrapper of a ``TensorType``.
        For example, you can use the following code to create a
        state type input that wraps a fp16 tensor with shape ``(2, 3)``::

            ct.StateType(
                wrapped_type=ct.TensorType(
                    shape=(2, 3),
                    dtype=np.float16
                ),
                name="state",
            )

        Parameters
        ----------
        wrapped_type: coremltools.converters.mil.input_types.InputType
            - The type wrapped in the state. 
            - Must be ``TensorType``.
              Note that the ``name`` and ``default_value`` of the wrapped ``TensorType`` must not be provided.

        name: str
            The name of the state.
            It must match the key of ``named_buffers()`` in the source TorchScript model.
        
coremltools/converters/mil/input_types.py