Description

First of all if this feature is already supported then please consider this as a question.

I'm trying to reproduce some results of existing algorithms such as SGD, AdamW, and also the recently proposed ScheduleFree. For the latter, in particular, there is an .eval() step that needs to be done before the validation phase, and I haven't figured out how to do that properly (there doesn't seem to be any indication in their submission code).

Any help would be greatly appreciated! Thank you in advance!