amogh112:

The place where it gets stuck is when the get() function is called in MLMAccuracy when writing to tensorboard and logger. Somehow it is not able to do all_reduce here for sum_metric  for this particular metric when more number of GPUs are there.


Can you verify that this allreduce gets called on all ranks? One possibility could be that some ranks don’t invoke this allreduce which would result in a freeze.