From e95e8279ef7a642797c3e5fedac47c7a65358a3f Mon Sep 17 00:00:00 2001
From: Vedant Puri <vedantpuri@gmail.com>
Date: Mon, 13 May 2024 20:22:31 -0400
Subject: [PATCH] configure BLAS

---
 examples/eg1.jl | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/examples/eg1.jl b/examples/eg1.jl
index 12edda2..39c3ef3 100644
--- 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()