MPI backend mnist gpu example error: "No space left on device"

[jwwandy]

pytorch-operator/examples/ddp/mnist/gpu/mnist_ddp_gpu.py

Lines 142 to 147 in 6538b24

	def average_gradients(model):
	""" Gradient averaging. """
	size = float(dist.get_world_size())
	group = dist.new_group([0])
	for param in model.parameters():
	dist.all_reduce(param.grad.data, op=dist.reduce_op.SUM, group=group)

When running MPI gpu examples with more epochs, average_gradients called dist.new_group in every backward iterations, which creates new ProcessGroupMPI each time and raises No space left on device error after few epocs as below. Also, original code creates group with group = dist.new_group([0]) instead of whole MPI world for back propagation.