DataParallel not using both GPUs equally

username: ilovrencic
cooked: <p>Hi everyone, I have a problem with my DataParallel model. I am wrapping my model with DataParallel, as I am trying to take advantage of 2 GPUs I have. However, when I check the GPUs performance when training the model, I notice that my model is using exclusively only one GPUs. I am not sure why is this happening, and I was wondering whether someone knows how to fix this issue.</p>
<pre><code class="lang-auto">class SummModel(nn.Module):

    def __init__(self, hidden_size=768):
        super(SummModel, self).__init__()

        self.hidden_size = hidden_size
        self.encoder = BertModel.from_pretrained('bert-base-uncased')

    """
    Method that represents one forward step in the model.
    
    Parameters:
        - text_id : list of tokenized texts
        - summary_id : list of tokenized summaries
        - candidate_id : list of list of N candidates
    """

    def forward(self, text_id, candidate_id, summary_id):
        batch_size = len(text_id)  # number of tokenized texts
        pad_id = 0  # pad symbol for Bert

        # document embedding
        mask = ~(text_id == pad_id)
        last_hidden_state = self.encoder(text_id, attention_mask=mask)[0]  # taking the state of the last layer in Bert model
        doc_emb = last_hidden_state[:, 0, :]
        assert doc_emb.size() == (batch_size, self.hidden_size)

        # summary embedding
        mask = ~(summary_id == pad_id)
        last_hidden_state = self.encoder(summary_id, attention_mask=mask)[0]
        sum_emb = last_hidden_state[:, 0, :]
        assert sum_emb.size() == (batch_size, self.hidden_size)

        # golden summary score
        summary_score = torch.cosine_similarity(sum_emb, doc_emb, dim=-1)

        # candidate embedding
        candidate_num = candidate_id.size(1)
        candidate_id = candidate_id.view(-1, candidate_id.size(-1))  # reshaping the tensor

        mask = ~(candidate_id == pad_id)
        last_hidden_state = self.encoder(candidate_id, attention_mask=mask)[0]
        candidate_emb = last_hidden_state[:, 0, :].view(batch_size, candidate_num, self.hidden_size)
        assert candidate_emb.size() == (batch_size, candidate_num, self.hidden_size)

        # calculate candidate scores
        doc_emb = doc_emb.unsqueeze(1).expand_as(candidate_emb)
        candidate_scores = torch.cosine_similarity(candidate_emb, doc_emb, dim=-1)
        assert candidate_scores.size() == (batch_size, candidate_num)

        return {"candidate_scores": candidate_scores, "summary_score": summary_score}
</code></pre>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/7/3/73328eae88d2f9307197e26ef2db4517c906cd97.png" data-download-href="https://discuss.pytorch.org/uploads/default/73328eae88d2f9307197e26ef2db4517c906cd97" title="Screenshot 2021-04-15 at 01.29.27"><img src="https://discuss.pytorch.org/uploads/default/original/3X/7/3/73328eae88d2f9307197e26ef2db4517c906cd97.png" alt="Screenshot 2021-04-15 at 01.29.27" data-base62-sha1="gr5e7nsKA0nSAxGGqjLmgqzoR6v" width="690" height="389" data-dominant-color="0E0E0E"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">Screenshot 2021-04-15 at 01.29.27</span><span class="informations">1112×628 31.4 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>EDIT: Fixed it by increasing batch size from 1 to 2. <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"></p>