diff --git a/irlc/ex02/dp.py b/irlc/ex02/dp.py
index e1cf1ea85f2b0f2d9c94732b17baab4b422293b1..24e061436debc4f0b07fe0cfea29fd88a49dfef0 100644
--- a/irlc/ex02/dp.py
+++ b/irlc/ex02/dp.py
@@ -7,18 +7,18 @@ References:
 from irlc.ex02.deterministic_inventory import DeterministicInventoryDPModel
 from irlc.ex02.dp_model import DPModel
 
-def DP_stochastic(model: DPModel): 
+def DP_stochastic(model: DPModel) -> tuple[list[dict], list[dict]]: 
     r"""
     Implement the stochastic DP algorithm. The implementation follows (Her25, Algorithm 1).
     Once you are done, you should be able to call the function as:
 
     .. runblock:: pycon
 
-        >>> from irlc.ex02.graph_traversal import SmallGraphDP
+        >>> from irlc.ex02.deterministic_inventory import DeterministicInventoryDPModel
         >>> from irlc.ex02.dp import DP_stochastic
-        >>> model = SmallGraphDP(t=5)  # Instantiate the small graph with target node 5
+        >>> model = DeterministicInventoryDPModel()  # Instantiate the deterministic DP model
         >>> J, pi = DP_stochastic(model)
-        >>> print(pi[0][2]) # Action taken in state ``x=2`` at time step ``k=0``.
+        >>> print(pi[0][2]) # Action taken in state ``x_0=2`` at time step ``k=0``.
 
     :param model: An instance of :class:`irlc.ex02.dp_model.DPModel` class. This represents the problem we wish to solve.
     :return: