actually, the code is a implementation of autodeeplab. I follow your instruction to modify my code:
class Architect_search(nn.Module):
    def __init__(self,archi_search,data_model):
        super(Architect_search, self).__init__()
        self.archi_model = nn.DataParallel(archi_search,device_ids=[0,1]).to('cuda:0')
        self.data_model = nn.DataParallel(data_model,device_ids=[2,3]).to('cuda:2')
    def forward(self, data_image,search_image ):
        search_result = self.archi_search(data_image)
        train_result = self.data_model(search_image)
        return search_result,train_result

bug got errors:
RuntimeError: module must have its parameters and buffers on device cuda:0 (device_ids[0]) but found one of them on device: cuda:2
specifically, the architect_search governs part of parameters of data_model.
class Architect_search (nn.Module) :
    def __init__(self, data_model, args):
        super(Architect_search, self).__init__()
        self.network_momentum = args.momentum
        self.network_weight_decay = args.weight_decay
        self.optimizer = torch.optim.Adam(model.arch_parameters(),
            lr=args.arch_lr, betas=(0.5, 0.999), weight_decay=args.arch_weight_decay)
        self.data_model = data_model


    def forward(self, input_valid):

            logits = self.model(input_valid)

            return logits

data_model is a common nn.Module class. arch_parameters will return [(20,8),(12,4,3)]
20 potential connection x 8 opperation types every block.  12 layers x4blockx3concatnation
WX20190716-111419@2x.png2108×486 116 KB