Skip to content

Commit

Permalink
gpu docstring
Browse files Browse the repository at this point in the history
  • Loading branch information
CarloLucibello committed Jul 16, 2023
1 parent 7462d4a commit 7f6e9d3
Showing 1 changed file with 6 additions and 2 deletions.
8 changes: 6 additions & 2 deletions src/functor.jl
Original file line number Diff line number Diff line change
Expand Up @@ -225,8 +225,10 @@ end
Copies `m` to the current GPU device (using current GPU backend), if one is available.
If no GPU is available, it does nothing (but prints a warning the first time).
On arrays, this calls CUDA's `cu`, which also changes arrays
with Float64 elements to Float32 while copying them to the device (same for AMDGPU).
When the backed is set to "CUDA", when called on arrays it calls `CUDA.cu`,
which also changes arrays with Float64 elements to Float32 while copying them to the device.
Similar conversions happen for "AMDGPU" and "Metal" backends.
To act on arrays within a struct, the struct type must be marked with [`@functor`](@ref).
Use [`cpu`](@ref) to copy back to ordinary `Array`s.
Expand All @@ -235,6 +237,8 @@ See also [`f32`](@ref) and [`f16`](@ref) to change element type only.
See the [CUDA.jl docs](https://juliagpu.github.io/CUDA.jl/stable/usage/multigpu/)
to help identify the current device.
See [`Flux.gpu_backend!`](@ref) for setting the backend.
# Example
```julia-repl
julia> m = Dense(rand(2, 3)) # constructed with Float64 weight matrix
Expand Down

0 comments on commit 7f6e9d3

Please sign in to comment.