Initialize CUDA lazily.

src/Flux.jl

@@ -30,7 +30,7 @@ export Descent, ADAM, Momentum, Nesterov, RMSProp,
 
 
 using CUDA
-const use_cuda = Ref(false)
+const use_cuda = Ref{Union{Nothing,Bool}}(nothing)
 
 include("utils.jl")
 include("zeros.jl")
@@ -57,13 +57,4 @@ include("deprecations.jl")
 
 include("cuda/cuda.jl")
 
-function __init__()
-  use_cuda[] = CUDA.functional() # Can be overridden after load with `Flux.use_cuda[] = false`
-  if CUDA.functional()
-    if !CUDA.has_cudnn()
-      @warn "CUDA.jl found cuda, but did not find libcudnn. Some functionality will not be available."
-    end
-  end
-end
-
 end # module

src/functor.jl

@@ -131,7 +131,16 @@ julia> typeof(m_gpu.W) # notice the type of the array changed to a CuArray
 CuArray{Float32, 2}
 ```
 """
-gpu(x) = use_cuda[] ? fmap(_gpu_array, x; exclude = _isbitsarray) : x
+function gpu(x)
+  if use_cuda[] === nothing
+    use_cuda[] = CUDA.functional()
+    if use_cuda[] && !CUDA.has_cudnn()
+      @warn "CUDA.jl found cuda, but did not find libcudnn. Some functionality will not be available."
+    end
+  end
+
+  use_cuda[] ? fmap(_gpu_array, x; exclude = _isbitsarray) : x
+end
 
 _gpu_array(x::AbstractArray) = CUDA.cu(x)
 

src/losses/Losses.jl

@@ -24,6 +24,5 @@ export mse, mae, msle,
 include("utils.jl")
 include("functions.jl")
 include("ctc.jl")
-if CUDA.functional() include("ctc-gpu.jl") end
 
 end #module