From b81560607b0d8f761df46b0f87dc9a496811dee8 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Eren=20G=C3=B6lge?= <egolge@coqui.ai>
Date: Thu, 22 Jul 2021 14:20:43 +0200
Subject: [PATCH] Update docstrings

---
 TTS/tts/layers/glow_tts/monotonic_align/__init__.py | 6 ++++--
 TTS/tts/layers/losses.py                            | 6 +++---
 2 files changed, 7 insertions(+), 5 deletions(-)

diff --git a/TTS/tts/layers/glow_tts/monotonic_align/__init__.py b/TTS/tts/layers/glow_tts/monotonic_align/__init__.py
index 5cbfd8fc..7757ecf8 100644
--- a/TTS/tts/layers/glow_tts/monotonic_align/__init__.py
+++ b/TTS/tts/layers/glow_tts/monotonic_align/__init__.py
@@ -21,8 +21,10 @@ def convert_pad_shape(pad_shape):
 
 def generate_path(duration, mask):
     """
-    duration: [b, t_x]
-    mask: [b, t_x, t_y]
+    Shapes:
+        - duration: :math:`[B, T_en]`
+        - mask: :math:'[B, T_en, T_de]`
+        - path: :math:`[B, T_en, T_de]`
     """
     device = duration.device
     b, t_x, t_y = mask.shape
diff --git a/TTS/tts/layers/losses.py b/TTS/tts/layers/losses.py
index 71e7e4fc..efe64e2b 100644
--- a/TTS/tts/layers/losses.py
+++ b/TTS/tts/layers/losses.py
@@ -69,9 +69,9 @@ class MSELossMasked(nn.Module):
             length: A Variable containing a LongTensor of size (batch,)
                 which contains the length of each data in a batch.
         Shapes:
-            x: B x T X D
-            target: B x T x D
-            length: B
+            - x: :math:`[B, T, D]`
+            - target: :math:`[B, T, D]`
+            - length: :math:`B`
         Returns:
             loss: An average loss value in range [0, 1] masked by the length.
         """