reintegrate has_cudnn check (#2200)

* reintegrate has_cudnn check

* typo

Author: CarloLucibello

Project.toml

@@ -23,6 +23,7 @@ SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
+cuDNN = "02a925ec-e4fe-4b08-9a7e-0d78e3d38ccd"
 
 [compat]
 AMDGPU = "0.4.8"
@@ -42,6 +43,7 @@ Reexport = "0.2, 1.0"
 SpecialFunctions = "1.8.2, 2.1.2"
 StatsBase = "0.33"
 Zygote = "0.6.49"
+cuDNN = "1"
 julia = "1.6"
 
 [extensions]

src/Flux.jl

@@ -44,6 +44,7 @@ using .Train
 using .Train: setup
 
 using CUDA
+import cuDNN
 const use_cuda = Ref{Union{Nothing,Bool}}(nothing)
 
 using Adapt, Functors, OneHotArrays

src/functor.jl

@@ -253,12 +253,16 @@ end
 function check_use_cuda()
   if use_cuda[] === nothing
     use_cuda[] = CUDA.functional()
+    if use_cuda[] && !cuDNN.has_cudnn()
+      @warn "CUDA.jl found cuda, but did not find libcudnn. Some functionality will not be available."  maxlog=1
+    end
     if !(use_cuda[])
       @info """The GPU function is being called but the GPU is not accessible. 
                Defaulting back to the CPU. (No action is required if you want to run on the CPU).""" maxlog=1
     end
   end
 end
+
 ChainRulesCore.@non_differentiable check_use_cuda()
 
 # Precision