Thanks for ur kindly reply.
I finally find the problem, let me rewrite my codes:
>>> CUDA_VISIBLE_DEVICEs=1,3,4 python train.py

class model_A(nn.Module):
...
   def forward(self, input):
      delta = ......
      delta = delta.cuda()

class model_B(nn.Module):
...

class model_C(nn.Module):
   def __init__(self, .....):
      self.M_A = model_A(...)
      self.M_B = model_B(...)

   def train_model(...)
      # DataLoader
      ...
      for input in DataLoader:
           image, .... = input
           C = self.M_A(image)
           out = self.M_B(C)
           loss = criterion(out, label)
           ...

# Train Process:
device = torch.device("cuda:0")
model.load_weights(model_path)
model = torch.nn.DataParallel(model)
model.to(device)
model.module.train_model(...)

No matter how I set my GPU, it still kept train on my first GPU.
The reason is because I write a train_model instead of the forward method and train model by calling model.module.train_model.
It’s ok to train by calling model.train_model to train on one GPU; however, when training multi-GPU, model.module.train_model(model is wrapped in nn.DataParallel) won’t access the forward method in DataParallel which cause the model not be replicated to different devices and input is not scattered.
To solve this, I write the forward method and write the train process outside the model_C class
Still is there any way to train model by something like model.train_model on multi-GPU???
Many thanks!!