Autograd notebook inconsistencies

[yaugenst-flex]

There are some inconsistencies w.r.t. variable naming / function setup / optimization loop etc. in the autograd notebooks which would be good to address.