Summary:
1. Add "set_device" call to keep things consistent between all DDP
tutorials.
This was inspired by the following change in the PyTorch repo:
pytorch/examples#1285 (review)
2. Fix up the tutorial and add additional prints when the model exits.

Test Plan:
Ran tutorial with the applied changes and we see:
"""
Running basic DDP example on rank 3.
Running basic DDP example on rank 1.
Running basic DDP example on rank 2.
Running basic DDP example on rank 0.
Finished running basic DDP example on rank 0.
Finished running basic DDP example on rank 1.
Finished running basic DDP example on rank 3.
Finished running basic DDP example on rank 2.
Running DDP checkpoint example on rank 2.
Running DDP checkpoint example on rank 1.
Running DDP checkpoint example on rank 0.
Running DDP checkpoint example on rank 3.
Finished DDP checkpoint example on rank 0.
Finished DDP checkpoint example on rank 3.
Finished DDP checkpoint example on rank 1.
Finished DDP checkpoint example on rank 2.
Running DDP with model parallel example on rank 0.
Running DDP with model parallel example on rank 1.
Finished running DDP with model parallel example on rank 0.
Finished running DDP with model parallel example on rank 1.
"""