Enable running pytorch.torchvision.train with distributed data parallel

[ohaijen]

Two main changes to enable running the following command in DistributedDataParallel mode:

CUDA_VISIBLE_DEVICES=<GPUs> python -m torch.distributed.launch \
 --nproc_per_node <NUM GPUs> \
sparseml.torchvision.train  <TRAIN.PY ARGUMENTS>

Before, it would throw an error (if run as above) or run in DataParallel mode (if run from CLI). After this change, running from CLI will still run DP, but running as above will run error-free with DDP, which is generally the desired outcome.

Concretely, the following changes were necessary:

• Accept local_rank argument set by torch.distributed.launch
• in case of distributed/DDP, ensure model gets sent to single GPU

Additionally, update documentation to explain how to run with DDP.

[bfineran]

LGTM

[bfineran]

nice - looks like it will make a lot of our defaults work out of the box