How to get average loss across all ranks using custom loss function

Below is a simple training code

for i, (inputs, labels) in enumerate(trainloader):  
        inputs = inputs.to(device)  
        labels = labels.to(device)  
        pred = model(inputs)  
        loss = criterion(pred, labels)  
        optimizer.zero_grad()  
        model.backward(loss)  
        optimizer.step()  

In this code, although the gradient descent of the model can be performed automatically.
But the loss is different in each rank. So how do I get the average loss under all ranks?
Thank you.

Start with the provided training loop, especially criterion, model.backward(loss), and the per-rank loss behavior. Read the project's distributed-training guidance and identify how a custom loss can be aggregated across ranks. Done means the documentation explains how to obtain the average loss and clarifies whether the reported value is per-rank or global.