From d12fb98f2a1530a0c5ff46822c7d16c027d99d9c Mon Sep 17 00:00:00 2001
From: Mike Innes <mike.j.innes@gmail.com>
Date: Sun, 15 Apr 2018 20:29:25 +0100
Subject: [PATCH] nicer batchnorm shape error

---
 src/layers/normalise.jl | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/src/layers/normalise.jl b/src/layers/normalise.jl
index 88c06673..c0d4aabf 100644
--- a/src/layers/normalise.jl
+++ b/src/layers/normalise.jl
@@ -114,6 +114,8 @@ BatchNorm(chs::Integer, λ = identity;
             zeros(chs), ones(chs), ϵ, momentum, true)
 
 function (BN::BatchNorm)(x)
+  size(x, ndims(x)-1) == length(BN.β) ||
+    error("BatchNorm expected $(length(BN.β)) channels, got $(size(x, ndims(x)-1))")
   λ, γ, β = BN.λ, BN.γ, BN.β
   dims = length(size(x))
   channels = size(x, dims-1)