-record(adam_state, {
    m :: viva_tensor@tensor:tensor(),
    v :: viva_tensor@tensor:tensor(),
    step :: integer()
}).