Adam optimizer .eval() not working well with ScheduleFree

[bhavnicksm]

[bhavnicksm]

I plan to refactor the project so that the initialisation of variables happens locally, where the variables are expected to be used. That means there's less confusion and checking involved so makes things easier to understand even if it makes it harder to get overall memory usage at a glance with the init_state method.