add doc

rwth-i6 · Jan 17, 2025 · 9211382 · 9211382
1 parent 88b7322
commit 9211382
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/i6_models/parts/samplers/log_uniform.py b/i6_models/parts/samplers/log_uniform.py
@@ -37,7 +37,7 @@ def sample(self, num_samples: int) -> torch.Tensor:
         Returns a random tensor in the size of [num_samples].
 
         :param num_samples: number of samples.
-        :return:
+        :return: [num_samples]
         """
         return self._cat_sampler.sample(torch.Size([num_samples]))
 
@@ -46,6 +46,6 @@ def log_prob(self, indices: torch.Tensor) -> torch.Tensor:
         Return log-probability of the given indices in the size of [B x T]
 
         :param indices: the ground truth target labels as indices.
-        :return:
+        :return: [B x T]
         """
         return self._cat_sampler.log_prob(indices)