fixed missing argument and refactoring

[Ankur-singh]

The split_spec parameter is not being passed to the pipeline function, resulting in the entire model being treated as a single stage. This results in RuntimeError: Pipeline group size 2 cannot be larger than number of stages 1, as the code attempts to call build_stage multiple times with different devices.

Additionally, I have made some cosmetic changes to improve readability. Since this example is directly referenced in the official documentation and will be referred by many users, it is crucial that it works.