Resume from checkpoint broken with distributed optimizer-in-backward

[ebsmothers]

To repro: patch the changes in #2359 to enable optimizer-in-backward for two recipe tests that test resume from checkpoint functionality.

The single-device command succeeds:

pytest -m integration_test tests/recipes/test_full_finetune_single_device.py -k 'test_training_state_on_resume'

But the distributed one fails:

pytest -m integration_test tests/recipes/test_full_finetune_distributed.py -k 'test_training_state_on_resume[llama3/8B_full-llama3-tune-4-1-True]'

Full stack trace of the failure can be seen here.