Note that for accumulating the loss in each rank, I use this class:
class AvgMeter:
    def __init__(self):
        self.reset()

    def reset(self):
        self.sum = 0
        self.count = 0
        self.avg = 0

    def update(self, val, n=1):
        self.sum += val * n
        self.count += n
        self.avg = self.sum / self.count

In your code, you can use this class something like this:
losses = AvgMeter()
for batch in pbar:
    # load image and mask into device memory
    image = batch['image'].cuda(rank, non_blocking=True)
    mask = batch['mask'].cuda(rank, non_blocking=True)

    # pass images into model
    pred = model(image)

    # get loss
    loss = criteria(pred, mask)
    losses.update(loss.item(), image.size(0))  # accumulate the loss

    # update the model
    optimizer.zero_grad(set_to_none=True)
    loss.backward()
    optimizer.step()

# after each epoch
loss = losses.avg
global_loss = global_meters_all_avg(args, loss) 

The global_loss is the one all-reduced (averaged) across the ranks.