How to use multiple GPUs in customized network layers?

username: amyxwgu
cooked: <p>I built a network using customized layers. It runs fine on a single GPU but crashes when using two GPUs of a server. The codes and error messages are shown below. It seems that one of the tensors was split into the 2 GPUs while the other was not. Was it caused by the customized forward function? How should I solve it? Thanks!</p>
<p>Code:</p>
<pre><code class="lang-auto">os.environ['CUDA_VISIBLE_DEVICES'] = '0,1'

class Cov1(layers):

    def __init__(self, in_dim=Fdim, out_dim=Fdim, bias=True):
        super(Cov1, self).__init__(in_dim, out_dim, bias)

    def forward(self,seq,sum_idx):
        simcov1 = torch.zeros(seq.shape).cuda()

        for i in range(0,self.in_dim):
            SeqDist = Vsets(seq[:,i].unsqueeze(1))
            simcov1[:, i] = torch.mean(SeqDist * sum_idx, 1)

        simcov1 = 1 - simcov1
        if self.bias is not None:
            mean_dist = simcov1.matmul(self.weight) + self.bias
            return (mean_dist)
        else:
            mean_dist = simcov1.matmul(self.weight)
            return (mean_dist)
</code></pre>
<p>Error screen:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/6/1/61fb8a5ceffb6dd0ef9b2d0e784a08692296fad9.png" data-download-href="https://discuss.pytorch.org/uploads/default/61fb8a5ceffb6dd0ef9b2d0e784a08692296fad9" title="image"><img src="https://discuss.pytorch.org/uploads/default/original/3X/6/1/61fb8a5ceffb6dd0ef9b2d0e784a08692296fad9.png" alt="image" data-base62-sha1="dYNeJrlZdhRxNNTxGeu42Rw20kN" width="690" height="31" data-dominant-color="1E1E1E"><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">image</span><span class="informations">867×40 2.62 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>

username: eqy
cooked: <p>If you are using <code>DataParallel</code> the assumption is all input tensors have the same dimension along the first (batch) dimension. Otherwise the splitting behavior becomes tricky to reason about.</p>
<p>What are the input shapes (and the meaning of the dimensions) being passed and is <code>DataParallel</code> being used?</p>

username: amyxwgu
cooked: <p>Thanks eqy. Yes, DataParallel is used for the model: <code>model = nn.DataParallel(model).cuda()</code>. And the dimension of the first tensor <code>SeqDist</code> is 446x446 and the second tensor <code>sum_idx</code> is 446. The multiplication <code>SeqDist * sum_idx</code> is to select the rows specified in <code>sum_idx</code> and calculate each row’s average.<br>
If the first dimension of the tensors changes, I don’t know how to make the multiplication work…</p>

username: eqy
cooked: <p>In this case, can you simply make this data parallel by doing something like making seqdist (N, 446, 446) and sum_idx (N, 446)?</p>

username: amyxwgu
cooked: <p>It works~ Thanks a lot~ I revised my coding to meet the splitting mechanism. <img src="https://discuss.pytorch.org/images/emoji/apple/v.png?v=9" title=":v:" class="emoji" alt=":v:"></p>