BN supports batchsize=1

1 files changed, 2 insertions, 2 deletions

diff --git a/BatchNormalization.lua b/BatchNormalization.lua
index 8dfc576..9e57998 100644
--- a/BatchNormalization.lua
+++ b/BatchNormalization.lua
@@ -130,7 +130,7 @@ function BN:updateOutput(input)
       self.running_var:cdata(),
       self.save_mean:cdata(),
       self.save_std:cdata(),
-      self.train,
+      self.train and (input:size(1) > 1), -- don't update running_[var,mean] when batchsize = 1
       self.momentum,
       self.eps)
 
@@ -162,7 +162,7 @@ local function backward(self, input, gradOutput, scale, gradInput, gradWeight, g
       self.running_var:cdata(),
       self.save_mean:cdata(),
       self.save_std:cdata(),
-      self.train,
+      self.train and (input:size(1) > 1), -- don't update running_[var,mean] when batchsize = 1
       scale,
       self.eps)