So I have been following the code and tutorial on using pytorch to do distributed machine learning here 1. I am able to run the code  (and it completes all the tasks) but my program does not terminate and I need to manually kill it using ctrl+C. The exact code is here 1
Right now after completing the task, it hangs after displaying the following warning messages

/anaconda3/lib/python3.7/site-packages/torch/distributed/distributed_c10d.py:406: UserWarning: For MPI backend, world_size (0) and rank (0) are ignored since they are assigned by the MPI runtime.
  "MPI runtime.".format(world_size, rank))
train_dist.py:72: UserWarning: Implicit dimension choice for log_softmax has been deprecated. Change the call to include dim=X as an argument.
  return F.log_softmax(x)
/anaconda3/lib/python3.7/site-packages/torch/distributed/distributed_c10d.py:125: UserWarning: torch.distributed.reduce_op is deprecated, please use torch.distributed.ReduceOp instead
  warnings.warn("torch.distributed.reduce_op is deprecated, please use "


I run the code using basic command

 /anaconda3/bin/mpirun -np 3 --host node-0,node-1,node-2 python train_dist.py


Do I need to add something in the code to exit gracefully?