diff --git a/dlordinal/losses/ordinal_ecoc_distance_loss.py b/dlordinal/losses/ordinal_ecoc_distance_loss.py
index b52d124..d1d1c7a 100644
--- a/dlordinal/losses/ordinal_ecoc_distance_loss.py
+++ b/dlordinal/losses/ordinal_ecoc_distance_loss.py
@@ -7,7 +7,7 @@
 
 class OrdinalECOCDistanceLoss(torch.nn.Module):
     """Ordinal ECOC distance loss from :footcite:t:`barbero2023error` for use
-    with :class:`dlordinal.models.OBDECOCModel`. Computes the MSE loss
+    with :class:`dlordinal.wrappers.OBDECOCModel`. Computes the MSE loss
     between the output of the model (class threshold probabilities) and
     the ideal output vector for each class.
 
diff --git a/dlordinal/output_layers/stick_breaking_layer.py b/dlordinal/output_layers/stick_breaking_layer.py
index f93b4dc..7f7ceaa 100644
--- a/dlordinal/output_layers/stick_breaking_layer.py
+++ b/dlordinal/output_layers/stick_breaking_layer.py
@@ -8,7 +8,7 @@ class StickBreakingLayer(Module):
     Parameters
     ----------
     input_shape: int
-        Input shape
+        Input shape, which refers to the number of neurons in the last fully connected layer
     num_classes: int
         Number of classes
     """
@@ -25,8 +25,8 @@ def get_stick_logits(self, x: torch.Tensor):
         x : torch.Tensor
             Input tensor
 
-        Returns:
-        --------
+        Returns
+        -------
         logits : torch.Tensor
             Logits of the stick breaking layer
         """
@@ -49,8 +49,8 @@ def forward(self, x) -> torch.Tensor:
         x : torch.Tensor
             Input tensor
 
-        Returns:
-        --------
+        Returns
+        -------
         logits : torch.Tensor
             Logits of the stick breaking layer
         """