configure BLAS

vpuri3 · May 14, 2024 · e95e827 · e95e827
1 parent 94696e5
commit e95e827
Showing 1 changed file with 7 additions and 0 deletions.
diff --git a/examples/eg1.jl b/examples/eg1.jl
@@ -13,6 +13,13 @@ using Zygote, Lux, ComponentArrays
 using LuxDeviceUtils, CUDA, LuxCUDA
 using BenchmarkTools
 
+# configure BLAS
+ncores = min(Sys.CPU_THREADS, length(Sys.cpu_info()))
+BLAS.set_num_threads(ncores)
+
+# configure CUDA
+CUDA.allowscalar(false)
+
 rng = Random.default_rng()
 Random.seed!(rng, 0)
 device = Lux.gpu_device()