From 3bc809f49e17e8463319920fbcef18725eea9d35 Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Sat, 10 Nov 2018 11:25:37 +0530 Subject: [PATCH] dropdims to make the array 2d --- src/cuda/cudnn.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/cuda/cudnn.jl b/src/cuda/cudnn.jl index 94424421..5d661889 100644 --- a/src/cuda/cudnn.jl +++ b/src/cuda/cudnn.jl @@ -45,8 +45,8 @@ batchnorm(g::CuArray{T}, b::CuArray{T}, x::CuArray{T, 2}, running_mean::CuArray{T}, running_var::CuArray{T}, momentum; cache = nothing, alpha = T(1), beta = T(0), eps = T(1e-5), training = true) where T<:Union{Float32, Float64} = - batchnorm(g, b, reshape(x, 1, 1, size(x, 1), size(x, 2)), running_mean, running_var, momentum, - cache = cache, alpha = alpha, beta = beta, eps = eps, training = training) + dropdims(batchnorm(g, b, reshape(x, 1, 1, size(x, 1), size(x, 2)), running_mean, running_var, momentum, + cache = cache, alpha = alpha, beta = beta, eps = eps, training = training), dims = (1, 2)) function batchnorm(g::CuArray{T}, b::CuArray{T}, x::Union{CuArray{T, 4},CuArray{T,5}}, running_mean::CuArray{T}, running_var::CuArray{T}, momentum;