add deivce in readme

vpuri3 · May 12, 2024 · 9185516 · 9185516
1 parent dc5226b
commit 9185516
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/README.md b/README.md
@@ -50,9 +50,9 @@ Chain(
 ```
 ```julia
 using CUDA, LuxCUDA
+device = Lux.gpu_device()
 
 x = rand32(rng, 1, 1000) |> device
-
 pM, stM = Lux.setup(rng, mlp) |> device
 pK, stK = Lux.setup(rng, kan) |> device