ToDevice

struct defined in module FluxTraining


			ToDevice(movedatafn, movemodelfn) <: Callback

Moves model and step data to a device using movedatafn for step data and movemodelfn for the model. For example ToDevice(Flux.gpu, Flux.gpu), moves them to a GPU if available. See ToGPU.

By default, only moves step.xs and step.ys, but this can be extended to other state by implementing on(::StepBegin, ::MyCustomPhase, ::ToDevice, learner).

Methods

There is 1 method for FluxTraining.ToDevice: