Set model device.
set_device(object, device, ...)
An neural network object.
(character
) Selected device.
Object of the same class with device set.
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