TastyPancakes/Flux.jl
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

cpu/gpu conveniences

Browse Source
This commit is contained in:
Mike J Innes 2018-02-26 23:10:59 +00:00
parent 15d1d3256b
commit 466b5c501a
2 changed files with 13 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
src/Flux.jl
View File

@ -10,7 +10,7 @@ using MacroTools: @forward
export Chain, Dense, RNN, LSTM, GRU, Conv, Conv2D, export Chain, Dense, RNN, LSTM, GRU, Conv, Conv2D,
Dropout, LayerNorm, BatchNorm, Dropout, LayerNorm, BatchNorm,
SGD, ADAM, Momentum, Nesterov, AMSGrad, SGD, ADAM, Momentum, Nesterov, AMSGrad,
param, params, mapleaves param, params, mapleaves, cpu, gpu
@reexport using NNlib @reexport using NNlib

12
src/treelike.jl
View File

@ -40,3 +40,15 @@ function params(m)
end end
params(m...) = params(m) params(m...) = params(m)
# CPU/GPU movement conveniences
cpu(x) = adapt(Array, x)
default_adaptor = Array
@require CuArrays begin
global default_adaptor = CuArray
end
gpu(x) = adapt(default_adaptor, x)