diff --git a/python/mxnet/visualization.py b/python/mxnet/visualization.py index 6e56dfb612ec..10a24241c730 100644 --- a/python/mxnet/visualization.py +++ b/python/mxnet/visualization.py @@ -240,6 +240,14 @@ def looks_like_weight(name): return True if name.endswith("_bias"): return True + if name.endswith("_var"): + return True + if name.endswith("_beta"): + return True + if name.endswith("_gamma"): + return True + if name.endswith("_mean"): + return True return False # make nodes