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.Dropout(self.dropout_rate),
            nn.Linear(vectors_size, vectors_size//2),
            nn.Tanh(),
            nn.Linear(vectors_size//2, features_size),
            nn.Tanh()
        )
        self.hidden = nn.Sequential(
            nn.Linear(features_size*2, hidden_size),
            nn.ReLU(),
        )

        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 = torch.cat([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