expose env var to disallow cuda

FluxML · Jan 14, 2022 · ebc1b30 · ebc1b30
1 parent 79dbbd6
commit ebc1b30
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/src/functor.jl b/src/functor.jl
@@ -179,6 +179,11 @@ end
 
 function check_use_cuda()
   if use_cuda[] === nothing
+    if get(ENV, "FLUX_DISALLOW_CUDA", "true") in ("1", "true")
+      use_cuda[] = false
+      @debug "Flux: CUDA disallowed" ENV["FLUX_DISALLOW_CUDA"]
+      return
+    end
     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."