Update initializer.py (#2163)

Make Xaiver(out, xxx) same to Torch
starimpact · May 16, 2016 · 9217027 · 9217027
1 parent a833098
commit 9217027
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/python/mxnet/initializer.py b/python/mxnet/initializer.py
@@ -235,8 +235,11 @@ def __init__(self, rnd_type="uniform", factor_type="avg", magnitude=3):
 
     def _init_weight(self, _, arr):
         shape = arr.shape
-        fan_in, fan_out = np.prod(shape[1:]), shape[0]
-        factor = 1
+        hw_scale = 1.
+        if len(shape) > 2:
+            hw_scale = np.prod(shape[2:])
+        fan_in, fan_out = shape[1] * hw_scale, shape[0] * hw_scale
+        factor = 1.
         if self.factor_type == "avg":
             factor = (fan_in + fan_out) / 2.0
         elif self.factor_type == "in":