diff --git a/src/diffusion_for_multi_scale_molecular_dynamics/utils/d3pm_utils.py b/src/diffusion_for_multi_scale_molecular_dynamics/utils/d3pm_utils.py
index 72075f6c..c9fb17fb 100644
--- a/src/diffusion_for_multi_scale_molecular_dynamics/utils/d3pm_utils.py
+++ b/src/diffusion_for_multi_scale_molecular_dynamics/utils/d3pm_utils.py
@@ -70,8 +70,9 @@ def get_probability_at_previous_time_step(
     small_epsilon: float,
     probability_at_zeroth_timestep_are_logits: bool = False,
 ) -> torch.Tensor:
-    r"""Compute :math:`P(a_{t-1} | a_t, \gamma_0)`, for given probability distribution :math:`\gamma_0` and a one-hot
-        distribution :math:`a_t`.
+    r"""Compute :math:`P(a_{t-1} | a_t, \gamma_0)`.
+
+    For given probability distribution :math:`\gamma_0` and a one-hot distribution :math:`a_t`.
 
     .. math::
         P(a_{t-1} | a_t, \gamma_0) = (\gamma_0^T \cdot \bar{Q}_{t-1} \cdot a_{t-1}) (a_{t-1}^T \cdot Q_t \cdot a_t) /