add forgotten reset for BN

1 files changed, 22 insertions, 20 deletions

diff --git a/test/test.lua b/test/test.lua
index 493ce53..10ceabb 100644
--- a/test/test.lua
+++ b/test/test.lua
@@ -1136,27 +1136,29 @@ end
 
 function cudnntest.SpatialBatchNormalization()
    -- batch
-      local h = math.random(5,10)
-      local w = math.random(5,10)
-      local bsz = math.random(1, 32)
-      local from = math.random(1, 32)
-      local input = torch.randn(bsz,from,h,w):cuda()
-      local gradOutput = torch.randn(bsz,from,h,w):cuda()
-      local cbn = cudnn.SpatialBatchNormalization(from, 1e-3):cuda()
-      local gbn = nn.SpatialBatchNormalization(from, 1e-3):cuda()
-      cbn.weight:copy(gbn.weight)
-      cbn.bias:copy(gbn.bias)
-      local rescuda = cbn:forward(input)
-      local groundtruth = gbn:forward(input)
-      local resgrad = cbn:backward(input, gradOutput)
-      local groundgrad = gbn:backward(input, gradOutput)
+   local h = math.random(5,10)
+   local w = math.random(5,10)
+   local bsz = math.random(1, 32)
+   local from = math.random(1, 32)
+   local input = torch.randn(bsz,from,h,w):cuda()
+   local gradOutput = torch.randn(bsz,from,h,w):cuda()
+   local cbn = cudnn.SpatialBatchNormalization(from, 1e-3):cuda()
+   local gbn = nn.SpatialBatchNormalization(from, 1e-3):cuda()
+   cbn.weight:copy(gbn.weight)
+   cbn.bias:copy(gbn.bias)
+   mytester:asserteq(cbn.running_mean:mean(), 0, 'error on BN running_mean init')
+   mytester:asserteq(cbn.running_std:mean(), 1, 'error on BN running_std init')
+   local rescuda = cbn:forward(input)
+   local groundtruth = gbn:forward(input)
+   local resgrad = cbn:backward(input, gradOutput)
+   local groundgrad = gbn:backward(input, gradOutput)
 
-      local error = rescuda:float() - groundtruth:float()
-      mytester:assertlt(error:abs():max(),
-                        precision_forward, 'error in batch normalization (forward) ')
-      error = resgrad:float() - groundgrad:float()
-      mytester:assertlt(error:abs():max(),
-                        precision_backward, 'error in batch normalization (backward) ')
+   local error = rescuda:float() - groundtruth:float()
+   mytester:assertlt(error:abs():max(),
+      precision_forward, 'error in batch normalization (forward) ')
+   error = resgrad:float() - groundgrad:float()
+   mytester:assertlt(error:abs():max(),
+      precision_backward, 'error in batch normalization (backward) ')
 end
 
 function cudnntest.SpatialCrossEntropyCriterion()