enhance one-hot documentation (#2521)

Signed-off-by: Wenqi Li <[email protected]>

Author: wyli

monai/networks/utils.py

@@ -37,13 +37,37 @@
 
 def one_hot(labels: torch.Tensor, num_classes: int, dtype: torch.dtype = torch.float, dim: int = 1) -> torch.Tensor:
     """
-    For a tensor `labels` of dimensions [B]1[spatial_dims], return a tensor of dimensions `[B]N[spatial_dims]`
-    for `num_classes` N number of classes.
+    For every value v in `labels`, the value in the output will be either 1 or 0. Each vector along the `dim`-th
+    dimension has the "one-hot" format, i.e., it has a total length of `num_classes`,
+    with a one and `num_class-1` zeros.
+    Note that this will include the background label, thus a binary mask should be treated as having two classes.
+
+    Args:
+        labels: input tensor of integers to be converted into the 'one-hot' format. Internally `labels` will be
+            converted into integers `labels.long()`.
+        num_classes: number of output channels, the corresponding length of `labels[dim]` will be converted to
+            `num_classes` from `1`.
+        dtype: the data type of the output one_hot label.
+        dim: the dimension to be converted to `num_classes` channels from `1` channel.
 
     Example:
 
-        For every value v = labels[b,1,h,w], the value in the result at [b,v,h,w] will be 1 and all others 0.
-        Note that this will include the background label, thus a binary mask should be treated as having 2 classes.
+    For a tensor `labels` of dimensions [B]1[spatial_dims], return a tensor of dimensions `[B]N[spatial_dims]`
+    when `num_classes=N` number of classes and `dim=1`.
+
+    .. code-block:: python
+
+        from monai.networks.utils import one_hot
+        import torch
+
+        a = torch.randint(0, 2, size=(1, 2, 2, 2))
+        out = one_hot(a, num_classes=2, dim=0)
+        print(out.shape)  # torch.Size([2, 2, 2, 2])
+
+        a = torch.randint(0, 2, size=(2, 1, 2, 2, 2))
+        out = one_hot(a, num_classes=2, dim=1)
+        print(out.shape)  # torch.Size([2, 2, 2, 2, 2])
+
     """
     if labels.dim() == 0:
         # if no channel dim, add it