Set model device.

set_device(object, device, ...)

Arguments

object

An neural network object.

device

(character) Selected device.

Value

Object of the same class with device set.

Examples

rnn_net <-
  model_rnn(
    input_size  = 1,
    output_size = 1,
    hidden_size = 10
  ) %>%
 set_device("cpu")

rnn_net
#> An `nn_module` containing 1,159 parameters.
#> 
#> ── Modules ─────────────────────────────────────────────────────────────────────
#> • initial_layer: <nn_nonlinear> #8 parameters
#> • rnn: <nn_gru> #1,140 parameters
#> • final_layer: <nn_linear> #11 parameters