When I use torch.nn.DataParallel(), the code run correctly.
Thank you so much!
Here is a minimal reproduce-able example:

  
      github.com
  
  
    lzkzls/DDPtest/blob/master/DDP_test.py 3
import torch, torchvision
import torch.nn as nn
import torch.distributed as dist
import torchvision.transforms as transforms
import torch.optim as optim


#input (1,28,28)
class Discriminator(nn.Module):
    def __init__(self):
        super(Discriminator, self).__init__()
        self.conv2 = nn.ModuleList()
        self.conv2.append(nn.Sequential(nn.Conv2d(1, 16, 3, stride=2, padding=1),
                                        nn.BatchNorm2d(16),
                                        nn.LeakyReLU(negative_slope=0.2)
        ))
        
        self.conv2.append(nn.Sequential(nn.Conv2d(16, 32, 3, stride=2, padding=1),
                                        nn.BatchNorm2d(32),
                                        nn.LeakyReLU(negative_slope=0.2)


  This file has been truncated. show original

  
  
    
    
  
  


and I run this code with command: sudo CUDA_VISIBLE_DEVICES=0,1,2,3 python3 -m torch.distributed.launch --nproc_per_node=4 DDP_test.py.
note: When I remove the comment of line 77 and 78, there will be an error:
image1348×160 8.39 KB