From 96ffcb2dc5fc768dd62dfb8149f0fdf9a6027544 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E2=80=9Cwenjuan=E2=80=9D?= Date: Thu, 2 Dec 2021 19:28:01 +0800 Subject: [PATCH] docs(mge/functional): update functional.batch_norm docstring --- imperative/python/megengine/functional/nn.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/imperative/python/megengine/functional/nn.py b/imperative/python/megengine/functional/nn.py index c127cc572..fd957db30 100644 --- a/imperative/python/megengine/functional/nn.py +++ b/imperative/python/megengine/functional/nn.py @@ -1194,6 +1194,13 @@ def batch_norm( eps: a value added to the denominator for numerical stability. Default: 1e-5 inplace: whether to update ``running_mean`` and ``running_var`` inplace or return new tensors. Default: True + compute_mode: When set to 'default', no special requirements will be + placed on the precision of intermediate results. When set to 'float32', + float32 would be used for accumulator and intermediate result, but only + effective when input and output are of float16 dtype. + param_dim: a value indicating in which format the parameters are. + Default: 'dim_1c11', which means NCHW format. + And 'dim_111c' means NHWC format. """ if inp.ndim != 4: raise NotImplementedError("batch_norm for ndim != 4") -- GitLab