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