Use a simple auto wrapping policy

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