I changed the ‘RandomLayer’ code and it’s working with 2 gpus, but still failed with 3 or 4 gpus. I donnot know why, so could you please have a loot at my question? In https://github.com/pytorch/pytorch/issues/50087,@ngimel told me that you cannot override .cuda() method this way, it’s not supported, you need to override _apply method instead. But I don’t know how to implemented it.
@ptrblck
Here is my modified RandomLyaer codes, because len(self.random_matrix) is actually 2 in my project, so I simply split ‘self.random_matrix[2]’ to 2 scalars. 2 gpus is working, 3 or more gpus will be failed.
class RandomLayer(nn.Module):
    def __init__(self, batch_size, 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_matrix0 = torch.nn.Parameter(torch.rand(input_dim_list[0], output_dim))
        self.random_matrix1 = torch.nn.Parameter(torch.rand(input_dim_list[1], output_dim))

    def forward(self, input_list):
        assert(len(input_list) == 2)
        return_list0 = torch.mm(input_list[0], self.random_matrix0)
        return_list1 = torch.mm(input_list[1], self.random_matrix1)
        return_tensor = return_list0 / math.pow(float(self.output_dim), 1.0 / 2)
        return_tensor = torch.mul(return_tensor, return_list1)
        return return_tensor