How to solve the different GPU devices when calculate the gradient of images?

username: FelixLe
cooked: <p><strong>I have 4 GPUs</strong>:</p>
<pre><code class="lang-auto">gpus = list(config.GPUS)
gradient_loss = GradientLoss
model = nn.DataParallel(model, device=gpus).cuda()
</code></pre>
<p><strong>Gradient loss is computed as below</strong>:</p>
<pre><code class="lang-auto">class GradientLoss(nn.Module):
    def __init__(self, channels=3):
        super(GradientLoss, self).__init__()

        pos = torch.from_numpy(np.identity(channels, dtype=np.float32))
        neg = -1 * pos
        self.filter_x = torch.stack((neg, pos)).unsqueeze(0).permute(3, 2, 0, 1).cuda()
        self.filter_y = torch.stack((pos.unsqueeze(0), neg.unsqueeze(0))).permute(3, 2, 0, 1).cuda()

    def forward(self, prediction, target):
        gen_frames_x = nn.functional.pad(prediction, [0, 1, 0, 0])
        gen_frames_y = nn.functional.pad(prediction, [0, 0, 0, 1])
        gt_frames_x = nn.functional.pad(target, [0, 1, 0, 0])
        gt_frames_y = nn.functional.pad(target, [0, 0, 0, 1])

        gen_dx = torch.abs(nn.functional.conv2d(gen_frames_x, self.filter_x))
        gen_dy = torch.abs(nn.functional.conv2d(gen_frames_y, self.filter_y))
        gt_dx = torch.abs(nn.functional.conv2d(gt_frames_x, self.filter_x))
        gt_dy = torch.abs(nn.functional.conv2d(gt_frames_y, self.filter_y))

        grad_diff_x = torch.abs(gt_dx - gen_dx)
        grad_diff_y = torch.abs(gt_dy - gen_dy)

        return torch.mean(grad_diff_x + grad_diff_y)
</code></pre>
<p><strong>The error</strong>:<br>
line 32, in forward<br>
gen_dx = torch.abs(nn.functional.conv2d(gen_frames_x, self.filter_x))<br>
RuntimeError: Expected tensor for argument <span class="hashtag">#1</span> ‘input’ to have the same device as tensor for argument <span class="hashtag">#2</span> ‘weight’; but device 1 does not equal 0 (while checking arguments for cudnn_convolution)</p>
<p>The reason is that the self.filter_x uses a GPU device (eg: cuda(0)) but the prediction and target images use another GPU device (eg: cuda(1)). How to solve this issue? Tks.</p>

username: comely420x
cooked: <p>Try changing the device of the first tensor using:</p>
<pre><code class="lang-auto">gen_dx = torch.abs(nn.Functional.conv2d((gen_frame_x).to('cuda:0'),self.filter_x))
</code></pre>
<p>Does this help?</p>

username: FelixLe
cooked: <p>Thanks for your advice.<br>
The prediction and target use all GPUs (4 GPUs). This line is suitable for the prediction and target of the first GPU.</p>
<pre><code class="lang-auto">gen_dx = torch.abs(nn.Functional.conv2d((gen_frame_x).to('cuda:0'),self.filter_x))
</code></pre>
<p>I need</p>
<pre><code class="lang-auto">gen_dx = torch.abs(nn.Functional.conv2d((gen_frame_x).to('cuda:0'),self.filter_x))
gen_dx = torch.abs(nn.Functional.conv2d((gen_frame_x).to('cuda:1'),self.filter_x))
gen_dx = torch.abs(nn.Functional.conv2d((gen_frame_x).to('cuda:2'),self.filter_x))
gen_dx = torch.abs(nn.Functional.conv2d((gen_frame_x).to('cuda:3'),self.filter_x))
</code></pre>
<p>How to do this?</p>

username: ptrblck
cooked: <p>Are you using the <code>GradientLoss</code> inside the <code>nn.DataParallel</code> module?<br>
If so, remove the <code>cuda()</code> operations from the <code>GradientLoss.__init__</code> method, register the tensors as buffers via <code>torch.register_buffer</code>, and <code>nn.DataParallel</code> should be able to clone this module to all passed GPUs.</p>

username: FelixLe
cooked: <p>I tried. It worked. Thank you very much, <a class="mention" href="/u/ptrblck">@ptrblck</a>!</p>