fix(mge/doc): fix batch_norm2d doc

GitOrigin-RevId: 769ddd367b7825f23a6343f13426e95aaae7f570

python_module/megengine/functional/nn.py

@@ -280,27 +280,20 @@ def batch_norm2d(
 ) -> Tensor:
     """Applies batch normalization to the input.
 
-    :type inp: Tensor
-    :param inp: The input tensor.
-    :type num_features: int
-    :param num_features: usually the :math:`C` from an input of size
-        :math:`(N, C, H, W)` or the highest ranked dimension of an input with
-        less than 4D.
-    :type eps: float
+    :param inp: input tensor.
+    :param running_mean: tensor to store running mean.
+    :param running_var: tensor to store running variance.
+    :param weight: scaling tensor in the learnable affine parameters.
+        See :math:`\gamma` in :class:`~.BatchNorm2d`
+    :param bias: bias tensor in the learnable affine parameters.
+        See :math:`\beta` in :class:`~.BatchNorm2d`
+    :param training: a boolean value to indicate whether batch norm is performed
+        in traning mode. Default: ``False``
+    :param momentum: the value used for the ``running_mean`` and ``running_var``
+        computation.
+        Default: 0.9
     :param eps: a value added to the denominator for numerical stability.
         Default: 1e-5.
-    :type momentum: float
-    :param momentum: the value used for the `running_mean` and `running_var`
-        computation.
-        Default: 0.1
-    :type affine: bool
-    :param affine: a boolean value that when set to ``True``, this module has
-        learnable affine parameters. Default: ``True``
-    :type track_running_stats: bool
-    :param track_running_stats: when set to ``True``, this module tracks the
-        running mean and variance. When set to ``False``, this module does not
-        track such statistics and always uses batch statistics in both training
-        and eval modes. Default: ``True``.
 
     Refer to :class:`~.BatchNorm2d` and :class:`~.BatchNorm1d` for more information.
     """