Does PyTorch work with MPI like other programs would? i.e., would the below work
$ mpiexec -n 4 python torch_script.py

when torch_script.py contains
torch.distributed.init_process_group(backend='mpi', world_size=4)