Follow-up improvements to RAFT training reference

[NicolasHug]

This issue is to keep track of the progress on some minor improvements related to the RAFT training reference, as a follow up to #4644

• support --weights param to support the prototype API, once pre-trained weights are available
• support --model instead of the --small flag
• support --epochs param instead of --num-steps, for consistency with the rest of the training references.
• save the scheduler and the optimizer, not just the model.
• try to keep function names as close as possible to existing references. For example validate() -> evaluate()
• Support --device (?? - not sure how useful it would be)
• Better support of weight serialization (move on CPU, avoid having the DDP model weights etc)

cc @datumbox

[NicolasHug]

I beleieve all the remaining points have been addressed by @YosuaMichael in #5590. Thanks Yosua!