Size dismatch when use multi-cuda(torch.nn.DataParallel)

Hi~ I’m using a custom model like this

class SimpleNN(nn.Module):
    def __init__(self, vectors_size, features_size, hidden_size=15, dropout_rate=0.1):
        super(SimpleNN, self).__init__()

        self.vectors_size = vectors_size
        self.features_size = features_size
        self.hidden_size = hidden_size
        self.dropout_rate = dropout_rate

        self.vectors_hidden = nn.Sequential(
            nn.Linear(vectors_size, vectors_size//2),
            nn.Linear(vectors_size//2, features_size),
        self.hidden = nn.Sequential(
            nn.Linear(features_size*2, hidden_size),

        self.output = nn.Linear(hidden_size, 2)

    def forward(self, pairs, features):
        features: (n_samples, features_size)
        vectors = pairs2vectors(train_pub, pairs).to(device)
        embedding_features = self.vectors_hidden(vectors)
        combined_features =[features, embedding_features], dim=1)
        return self.output(self.hidden(combined_features))

This model works well when i use only one cuda, but after ‘DataParallel’ used like below, it always tell me the size of features and embedding_features are not match, i find that the n_samples shape of features doesn’t follow my expectation just like another batch data, i dont know why and how to solve this problem.

if torch.cuda.device_count() > 1:
        print("Let's use", torch.cuda.device_count(), "GPUs!")
        model = nn.DataParallel(model)

BTW, here is the pic of error message is Error message