Merge pull request #40 from FluxML/DhairyaLGandhi-patch-1

Use new `PoolDims` constructor for NNlib v0.8

Author: DhairyaLGandhi

ext/NNlibCUDA/Project.toml

@@ -1,6 +1,6 @@
 name = "NNlibCUDA"
 uuid = "a00861dc-f156-4864-bf3c-e6376f28a68d"
-version = "0.1.11"
+version = "0.2.0"
 
 [deps]
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
@@ -11,7 +11,7 @@ Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 
 [compat]
 CUDA = "3.3.1"
-NNlib = "0.7.31"
+NNlib = "0.8"
 julia = "1.6"
 
 [extras]

ext/NNlibCUDA/src/cudnn/pooling.jl

@@ -41,8 +41,14 @@ end
 
 add1d(x) = reshape(x, 1, size(x)...)
 
-fix_pooldims_1d(pdims::PoolDims{1,K,S,P,D}) where {K,S,P,D} =
-        PoolDims{2,(1,K...),(1,S...),(0,0,P...),(1,D...)}((1,pdims.I...), pdims.C_in)
+function fix_pooldims_1d(pdims::PoolDims{1,K,S,P,D}) where {K,S,P,D}
+    PoolDims{2, K + 1, S + 1, P + 2, D + 1}((1, NNlib.input_size(pdims)...),
+                                            (1, NNlib.kernel_size(pdims)...),
+                                            NNlib.channels_in(pdims),
+                                            (1, NNlib.stride(pdims)...),
+                                            (0, 0, NNlib.padding(pdims)...),
+                                            (1, NNlib.dilation(pdims)...))
+end
 
 function maxpool!(y::DenseCuArray{T,3}, x::DenseCuArray{T,3}, pdims::PoolDims) where T<:CUDNNFloat
     maxpool!(add1d(y), add1d(x), fix_pooldims_1d(pdims))