chore: load LuxCUDA for correct gpu device

vpuri3 · Nov 14, 2024 · 0c24bc9 · 0c24bc9
1 parent 40e2a24
commit 0c24bc9
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/README.md b/README.md
@@ -31,7 +31,7 @@ y = layer(x, p, st)
 We compare the performance of different implementation of KAN with an MLP that has the same number of parameters (see `examples/eg1.jl`).
 ```julia
 using Lux, KolmogorovArnold
-using CUDA
+using LuxCUDA
 
 CUDA.allowscalar(false)
 device = Lux.gpu_device()