Hey @rzhang63 sorry that I mis-read your log picture. Are you using Windows? Currently PT Distributed does not support windows yet. We have a poll here: https://github.com/pytorch/pytorch/issues/37068
If running on Linux, it should work with a minor fix.
Change
labels = torch.randn(20, 5).to(rank)

to
labels = torch.randn(20, 10).to(rank)