Fix for getindex error in BatchNorm

2019-07-12 22:19:41 +05:30 · 2019-07-12 22:19:41 +05:30 · 2816fbb9b2
commit 2816fbb9b2
parent fc1c0d58ed
1 changed files with 1 additions and 1 deletions
--- a/src/layers/normalise.jl
+++ b/src/layers/normalise.jl
@ -136,7 +136,7 @@ function (BN::BatchNorm)(x)
  dims = length(size(x))
  channels = size(x, dims-1)
  affine_shape = ntuple(i->i == ndims(x) - 1 ? size(x, i) : 1, ndims(x))
-  m = prod(size(x)[1:end-2]) * size(x)[end]
+  m = trunc(Int, prod(size(x))/channels)
  γ = reshape(BN.γ, affine_shape...)
  β = reshape(BN.β, affine_shape...)
  if !istraining()