updating fsdp docs as per resolution to issue https://github.com/huggingface/accelerate/issues/1054 (#3)

- updating fsdp docs as per resolution to issue https://github.com/huggingface/accelerate/issues/1054 (48d3c971c8a350515cfb5300ee84457740bbcb25)


Co-authored-by: Sourab Mangrulkar <smangrul@users.noreply.huggingface.co>

code_samples/training_configuration/pytorch_fsdp

@@ -27,16 +27,27 @@ use_cpu: false
 <pre>
 from accelerate import Accelerator
 
-accelerator = Accelerator()
-- model, optimizer, dataloader, scheduler = accelerator.prepare(
--  model, optimizer, dataloader, scheduler
--)
-+model = accelerator.prepare(model)
-+# Optimizer can be any PyTorch optimizer class
-+optimizer = torch.optim.AdamW(params=model.parameters(), lr=lr)
-+optimizer, dataloader, scheduler = accelerator.prepare(
-+  optimizer, dataloader, scheduler
-+)
+def main():
+    accelerator = Accelerator()
+-    model, optimizer, dataloader, scheduler = accelerator.prepare(
+-      model, optimizer, dataloader, scheduler
+-    )
++    model = accelerator.prepare(model)
++    # Optimizer can be any PyTorch optimizer class
++    optimizer = torch.optim.AdamW(params=model.parameters(), lr=lr)
++    optimizer, dataloader, scheduler = accelerator.prepare(
++      optimizer, dataloader, scheduler
++    )
+    
+    ...
+
+    accelerator.unwrap_model(model).save_pretrained(
+            args.output_dir,
+            is_main_process=accelerator.is_main_process,
+            save_function=accelerator.save,
++            state_dict=accelerator.get_state_dict(model)
+    )
+    ...
 </pre>
 ##
 If the YAML was generated through the `accelerate config` command: