TastyPancakes/Flux.jl
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Make cuDNN take only 4D arrays

Browse Source
This commit is contained in:
Avik Pal 2018-06-20 15:54:38 +05:30
parent 3339ad5181
commit deb4950261
1 changed files with 4 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
src/cuda/cudnn.jl
View File

@ -37,7 +37,10 @@ end
bncache() = bncache(nothing, nothing) 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; running_mean::CuArray{T}, running_var::CuArray{T}, momentum;
cache = nothing, alpha = T(1), beta = T(0), cache = nothing, alpha = T(1), beta = T(0),
eps = T(1e-5), training = true) where T<:Union{Float32, Float64} eps = T(1e-5), training = true) where T<:Union{Float32, Float64}