Flux assumes only trainable parameters should be moved to GPU

ericphanson · August 15, 2021, 10:50am

I think you can specify which parameters are trainable separately from which parameters are the layers “children”. Eg for BatchNorms: https://github.com/FluxML/Flux.jl/blob/29a96b961badea84a3bb323257c1374ffebca2e7/src/layers/normalise.jl#L265.

Topic		Replies	Views
Implementing a custom layer Machine Learning	2	900	May 6, 2020
Apply boolean mask in loss function Machine Learning question , gpu , flux	1	959	July 23, 2020
Custom layer with new parameters in Flux.jl Machine Learning flux	3	1570	December 11, 2019
Flux.jl: training fails at GPU but works on CPU Machine Learning gpu , flux	1	647	September 19, 2019
Data Science lessons: Making "10 - Neural Networks" run on GPU? New to Julia gpu , flux	4	766	January 14, 2022