From c0994c77fc211e52415ed74eca4ca806e3382a9e Mon Sep 17 00:00:00 2001
From: Michael Abbott <32575566+mcabbott@users.noreply.github.com>
Date: Sat, 26 Nov 2022 14:34:47 -0500
Subject: [PATCH] maybe this example should run on the GPU, since it easily
 can, even though this is slower

---
 docs/src/models/quickstart.md | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/docs/src/models/quickstart.md b/docs/src/models/quickstart.md
index efdfde5f1e..a5d5bb9454 100644
--- a/docs/src/models/quickstart.md
+++ b/docs/src/models/quickstart.md
@@ -17,14 +17,14 @@ model = Chain(
     Dense(2 => 3, tanh),   # activation function inside layer
     BatchNorm(3),
     Dense(3 => 2),
-    softmax)
+    softmax) |> gpu        # move model to GPU, if available
 
 # The model encapsulates parameters, randomly initialised. Its initial output is:
-out1 = model(noisy)                                               # 2×1000 Matrix{Float32}
+out1 = model(noisy |> gpu) |> cpu                                 # 2×1000 Matrix{Float32}
 
 # To train the model, we use batches of 64 samples, and one-hot encoding:
 target = Flux.onehotbatch(truth, [true, false])                   # 2×1000 OneHotMatrix
-loader = Flux.DataLoader((noisy, target), batchsize=64, shuffle=true);
+loader = Flux.DataLoader((noisy, target) |> gpu, batchsize=64, shuffle=true);
 # 16-element DataLoader with first element: (2×64 Matrix{Float32}, 2×64 OneHotMatrix)
 
 pars = Flux.params(model)  # contains references to arrays in model
@@ -34,7 +34,7 @@ opt = Flux.Adam(0.01)      # will store optimiser momentum, etc.
 losses = []
 for epoch in 1:1_000
     for (x, y) in loader
-        loss, grad = withgradient(pars) do
+        loss, grad = Flux.withgradient(pars) do
             # Evaluate model and loss inside gradient context:
             y_hat = model(x)
             Flux.crossentropy(y_hat, y)
@@ -46,7 +46,7 @@ end
 
 pars  # parameters, momenta and output have all changed
 opt
-out2 = model(noisy)  # first row is prob. of true, second row p(false)
+out2 = model(noisy |> gpu) |> cpu  # first row is prob. of true, second row p(false)
 
 mean((out2[1,:] .> 0.5) .== truth)  # accuracy 94% so far!
 ```