fix batchnorm reset

stephenrawls · Oct 19, 2015 · eb13695 · eb13695
1 parent 89ad6c5
commit eb13695
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 0 deletions.
diff --git a/BatchNormalization.lua b/BatchNormalization.lua
@@ -59,6 +59,8 @@ end
 function BN:reset()
    self.weight:uniform()
    self.bias:zero()
+   self.running_mean:zero()
+   self.running_std:fill(1)
 end
 
 function BN:updateOutput(input)

diff --git a/SpatialBatchNormalization.lua b/SpatialBatchNormalization.lua
@@ -60,6 +60,8 @@ end
 function BN:reset()
    self.weight:uniform()
    self.bias:zero()
+   self.running_mean:zero()
+   self.running_std:fill(1)
 end
 
 function BN:updateOutput(input)