Fix Norm Layers, Again

src/cuda/cudnn.jl

@@ -1,14 +1,13 @@
 import NNlibCUDA: batchnorm, ∇batchnorm
 
-function (BN::Flux.BatchNorm)(x::Union{CuArray{T,2},CuArray{T,4},CuArray{T,5}},
-                              cache=nothing) where T<:Union{Float32, Float64}
+function (BN::Flux.BatchNorm)(x::CuArray{T},
+                              cache = nothing) where T<:Union{Float32, Float64}
 
-  @assert BN.affine "BatchNorm: only affine=true supported on gpu"
-  @assert BN.track_stats "BatchNorm: only track_stats=true supported on gpu"
-  @assert length(BN.β) == size(x, ndims(x)-1) "BatchNorm: input has wronng number of channels"
+  @assert BN.affine throw(ArgumentError("BatchNorm: only affine = true supported on gpu"))
+  @assert BN.track_stats throw(ArgumentError("BatchNorm: only track_stats = true supported on gpu"))
   return BN.λ.(batchnorm(BN.γ, BN.β, x, BN.μ, BN.σ², BN.momentum; 
-                  cache=cache, alpha=1, beta=0, eps=BN.ϵ, 
-                  training=Flux._isactive(BN)))
+                  cache = cache, alpha = 1, beta = 0, eps = BN.ϵ, 
+                  training = Flux._isactive(BN)))
 end
 
 @adjoint function batchnorm(g, b, x, running_mean, running_var, momentum; kw...)