Use a simple auto wrapping policy

This way, the example actually shows a way to properly distribute a
module using FSDP.

pytorch-fsdp-example/main.py

@@ -264,6 +264,11 @@ def main():
     model = fsdp.FullyShardedDataParallel(
         model,
         device_id=local_rank,
+        auto_wrap_policy=functools.partial(
+            fsdp.wrap.size_based_auto_wrap_policy,
+            # Wrap every 1B parameters.
+            min_num_params=int(1e9),
+        ),
     )
     loss_func = torch.nn.CrossEntropyLoss()