add ckpt for outer schedler

Author: samsja

open_diloco/train_fsdp.py

@@ -380,6 +380,7 @@ def outer_scheduler_fn(opt):
                 model=model,
                 optimizer=optimizer.inner_optimizer,
                 scheduler=scheduler,
+                outer_scheduler=optimizer.outer_scheduler,
                 outer_optimizer=optimizer.state_averager.optimizer,
                 scaler=scaler,
                 data_loader=train_dataloader,