-record(training_state, {
    tape :: viva_tensor@nn@autograd:tape(),
    layer1 :: viva_tensor@nn@layers:linear(),
    layer2 :: viva_tensor@nn@layers:linear()
}).