From 91855163816cedc30103352d665c5517f617bf18 Mon Sep 17 00:00:00 2001
From: Vedant Puri <vedantpuri@gmail.com>
Date: Sun, 12 May 2024 17:23:12 -0400
Subject: [PATCH] add deivce in readme

---
 README.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/README.md b/README.md
index 34f4e9f..d2ef229 100644
--- 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