diff --git a/dynamax/hidden_markov_model/models/abstractions.py b/dynamax/hidden_markov_model/models/abstractions.py
index defc70762..1dd6b946b 100644
--- a/dynamax/hidden_markov_model/models/abstractions.py
+++ b/dynamax/hidden_markov_model/models/abstractions.py
@@ -295,7 +295,7 @@ def collect_suff_stats(self,
             PyTree of sufficient statistics for updating the transition distribution
 
         """
-        return posterior.trans_probs, pytree_slice(inputs, slice(1, None))
+        return posterior.trans_probs, inputs 
 
     def initialize_m_step_state(self, params: ParameterSet, props:PropertySet) -> Any:
         """Initialize any required state for the M step.
diff --git a/pyproject.toml b/pyproject.toml
index bd7e910fe..31891f388 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -70,7 +70,8 @@ test = [
     "coverage",
     "pytest>=3.9",
     "pytest-cov",
-    "interrogate>=1.5.0"
+    "interrogate>=1.5.0",
+    "ipython"
 ]
 
 dev = [
@@ -91,7 +92,8 @@ dev = [
     "coverage",
     "pytest>=3.9",
     "pytest-cov",
-    "interrogate>=1.5.0"
+    "interrogate>=1.5.0",
+    "ipython"
 ]
 
 [tool.setuptools.packages.find]