Add missing expected tensor shapes in monotonic regularization loss d…

…ocstrings (#183)
vitalab · Nov 8, 2023 · 0c1ab44 · 0c1ab44
1 parent af24988
commit 0c1ab44
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 4 deletions.
diff --git a/vital/metrics/train/functional.py b/vital/metrics/train/functional.py
@@ -121,8 +121,8 @@ def monotonic_regularization_loss(input: Tensor, target: Tensor, delta: float) -
           (link to the paper: https://arxiv.org/pdf/2004.05485.pdf)
 
     Args:
-        input: Input values to regularize so that they have a monotonic relationship with the `target` values.
-        target: Values used to determine the target monotonic ordering of the values.
+        input: (N, [1]), Input values to regularize so that they have a monotonic relationship with the `target` values.
+        target: (N, [1]), Values used to determine the target monotonic ordering of the values.
         delta: Hyperparameter that decides the spread of the posterior distribution.
 
     Returns:

diff --git a/vital/metrics/train/metric.py b/vital/metrics/train/metric.py
@@ -80,8 +80,9 @@ def forward(self, input: Tensor, target: Tensor) -> Tensor:
         """Actual metric calculation.
 
         Args:
-            input: Input values to regularize so that they have a monotonic relationship with the `target` values.
-            target: Values used to determine the target monotonic ordering of the values.
+            input: (N, [1]), Input values to regularize so that they have a monotonic relationship with the `target`
+                values.
+            target: (N, [1]), Values used to determine the target monotonic ordering of the values.
 
         Returns:
             (1,), Calculated monotonic regularization loss.