if loss.item() > 1.8: 
    loss.backward()
else:
    print("skipping batch:", loss.item())

The above might be the cause of the problem. When using DistributedDataParallel, backward() pass will trigger gradient synchronization communication (all_reduce) across all processes, meaning that all processes need to agree on the number and order of all_reduce calls. However, the above code seems to skip the backward pass in some process but not guarantee to skip in other processes? If that is the case, then processes could run in to desync and cause hang.