From deb495026184b67e407e7e31adb164f3bd9f00cd Mon Sep 17 00:00:00 2001
From: Avik Pal <avikpal@iitk.ac.in>
Date: Wed, 20 Jun 2018 15:54:38 +0530
Subject: [PATCH] Make cuDNN take only 4D arrays

---
 src/cuda/cudnn.jl | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/src/cuda/cudnn.jl b/src/cuda/cudnn.jl
index c7d997b9..9948ef37 100644
--- a/src/cuda/cudnn.jl
+++ b/src/cuda/cudnn.jl
@@ -37,7 +37,10 @@ end
 
 bncache() = bncache(nothing, nothing)
 
-function batchnorm(g::CuArray{T}, b::CuArray{T}, x::CuArray{T},
+# CuDNN supports only 4D and 5D Tensors for BatchNorm Operations
+# so use the native julia code when doing batchnorm on a 2D Array
+
+function batchnorm(g::CuArray{T}, b::CuArray{T}, x::CuArray{T, 4},
                    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}