# How choose which GPU to send the workload to in Flux.jl?

**URL:** <https://discourse.julialang.org/t/how-choose-which-gpu-to-send-the-workload-to-in-flux-jl/16337>\
**Category:** Machine Learning\
**Tags:** first-steps, flux\
**Created:** [October 15, 2018, 10:25am UTC](https://discourse.julialang.org/t/how-choose-which-gpu-to-send-the-workload-to-in-flux-jl/16337 "2018-10-15T10:25:50Z")\
**Posts on this page:** 4\
**Page:** 1

<div class="post-metadata">

**Author:** ![xiaodai](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/xiaodai/32/15937_2.png) [@xiaodai](https://discourse.julialang.org/u/xiaodai)\
**Post date:** [October 15, 2018, 10:25am UTC](https://discourse.julialang.org/t/how-choose-which-gpu-to-send-the-workload-to-in-flux-jl/16337/1 "2018-10-15T10:25:50Z")

</div>

I have set up a neural network in Flux and I would like to train it on a GPU. HOwever my laptop has two GPUs and one is an onboard GPU so I like to send the workload to the better GPU.

But how I choose which GPU to send the workload to?

```julia
using Flux

policy = Flux.Chain(
  Conv((2,2), 1=>128, relu)
  ,Conv((2,2), 128=>128, relu)
  ,x -> reshape(x, :, size(x,4))
  ,Dense(512, 512, relu)
  ,Dense(512, 4)
  ,softmax
  ) |> gpu

@time training_sample = cat((init_game() .|> Float64 for i=1:10_000)..., dims=4) |> gpu
@time policy(training_sample)

```

---

<div class="post-metadata">

**Author:** ![sdanisch](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/sdanisch/32/1406_2.png) [@sdanisch](https://discourse.julialang.org/u/sdanisch)\
**Post date:** [October 15, 2018, 10:30am UTC](https://discourse.julialang.org/t/how-choose-which-gpu-to-send-the-workload-to-in-flux-jl/16337/2 "2018-10-15T10:30:04Z")

</div>

If you’re using CUDA you don’t need to worry about that, since it won’t be able to use your onboard GPU 😉  
Anyways, this is how you could do it:

[from the CUDANative Docs](http://juliagpu.github.io/CUDAnative.jl/latest/man/usage.html#Quick-start-1):

> This code is executed in a default, global context for the first device in your system. Similar to `cudaSetDevice` , you can switch devices by calling CUDAnative’s `device!` function:

```julia
# change the active device
device!(1)

# the same, but only temporarily
device!(2) do
    # ...
end

```

If nothing weird is happening in CuArrays, this should also set the device for Flux using the GPU 😉

---

<div class="post-metadata">

**Author:** ![saikhu](https://avatars.discourse-cdn.com/v4/letter/s/65b543/32.png) [@saikhu](https://discourse.julialang.org/u/saikhu)\
**Post date:** [March 11, 2022, 8:01am UTC](https://discourse.julialang.org/t/how-choose-which-gpu-to-send-the-workload-to-in-flux-jl/16337/3 "2022-03-11T08:01:47Z")

</div>

can you please also tell me, how to choose the number of gpus, like I have 3 gpus but the flux automatically use only first one. I want to distribute the load on all three. Thanks

---

<div class="post-metadata">

**Author:** ![ToucheSir](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/touchesir/32/14411_2.png) [@ToucheSir](https://discourse.julialang.org/u/ToucheSir)\
**Post date:** [March 11, 2022, 4:14pm UTC](https://discourse.julialang.org/t/how-choose-which-gpu-to-send-the-workload-to-in-flux-jl/16337/4 "2022-03-11T16:14:15Z")

</div>

Simon’s code above just that? You’d have to change between all 3 devices.