Hi ,ptrblck. Thank you for your kindly reply. I try to use nn.ParameterList in multi-gpus mode in your coe snippet, but it still doesn’t work. Here is the source code named ‘dp.py’, single gpu is OK.
#!/usr/bin/env python
# -*- coding: utf-8 -*-

import torch
import torch.nn as nn
import math

feat_d = 16
feat_d2 = 8
feat_d3 = 8

class RandomLayer(nn.Module):
    def __init__(self, input_dim_list=[], output_dim=1024):
        super(RandomLayer, self).__init__()
        self.input_num = len(input_dim_list)# 2
        self.output_dim = output_dim
        self.random_matrix = torch.nn.ParameterList([torch.nn.Parameter(torch.rand(input_dim_list[i], output_dim)) for i in range(self.input_num)])
        print('*****')
        print(len(self.random_matrix))
        print(self.random_matrix)

    def forward(self, input_list):
        print('&&&&&')
        print(len(self.random_matrix))
        print(self.random_matrix)
        return_list = [torch.mm(input_list[i], self.random_matrix[i]) for i in range(self.input_num)]
        return_tensor = return_list[0] / math.pow(float(self.output_dim), 1.0 / 2)
        for single in return_list[1:]:
            return_tensor = torch.mul(return_tensor, single)
        return return_tensor


class MyModule(nn.Module):
    def __init__(self):
        super(MyModule, self).__init__()
        self.rand = RandomLayer([feat_d, feat_d2], feat_d3)

    def forward(self, x1, x2):
        print(x1.device)
        x = self.rand([x1, x2])
        return x

model = MyModule()
model = nn.DataParallel(model).cuda()
x1 = torch.randn(4, feat_d).cuda()
x2 = torch.randn(4, feat_d2).cuda()
out = model(x1, x2)
print(out.device)

when I running that code with the command ‘CUDA_VISIBLE_DEVICES=4,5 python dp.py’, ‘len(self.random_matrix)’ is 0. I don’t know why, so could you have a look at this?