Using grid_sample on multiple GPUs

username: sunshineatnoon
cooked: <p>Hi, I have a model with a <a href="https://pytorch.org/docs/stable/nn.functional.html?highlight=grid_sample#torch.nn.functional.grid_sample" rel="nofollow noopener">grid_sample</a> layer, I tried to train my model on multiple GPUs, but got the following error:</p>
<pre><code class="lang-auto">RuntimeError: grid_sampler(): expected input and grid to be on same device, but input is on cuda:1 and grid is on cuda:0
</code></pre>
<p>Is there anyway to use this layer on multiple GPUs? Thanks</p>

username: ptrblck
cooked: <p>The <code>input</code> and <code>grid</code> should be on the same device.<br>
If you are creating one of these tensors manually in the <code>forward</code> or pass it to the <code>forward</code> method, make sure to transfer it to the same device, e.g. by using:</p>
<pre><code class="lang-python">grid = grid.to(x.device)
</code></pre>

username: sunshineatnoon
cooked: <p>Thanks for the reply. I got a <code>segmentation default</code> when moving either of the grid or the input to the same device by <code>input = input.to(grid.get_device())</code>.</p>

username: sunshineatnoon
cooked: <p>My grid is actually the same for all inputs, so I stored it using <code>self.grid = grid</code> and using <code>grid_sample(input, self.grid)</code>. Do you think this causes the problem? But I think it’s inefficient to pass the grid every forward.</p>

username: ptrblck
cooked: <p>Try to use <code>grid.device</code> instead.</p>
<aside class="quote no-group" data-username="sunshineatnoon" data-post="4" data-topic="53405">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/sunshineatnoon/48/13446_2.png" class="avatar"> sunshineatnoon:</div>
<blockquote>
<p>But I think it’s inefficient to pass the grid every forward.</p>
</blockquote>
</aside>
<p>Might be and you should stick to your work flow, as I was just using it as an example. <img src="https://discuss.pytorch.org/images/emoji/apple/wink.png?v=12" title=":wink:" class="emoji" alt=":wink:" loading="lazy" width="20" height="20"></p>
<p>You could also try to register <code>self.grid</code> as a <code>buffer</code> using <code>self.register_buffer</code>, which would move the tensor automatically using <code>model.to()</code>.</p>

username: sunshineatnoon
cooked: <p><code>register_buffer</code> solves my problem. The <code>segmentation default</code> actually comes from other parts. It seems when training on multiple GPUs, we cannot call <code>.cuda()</code> during the forward path, so everything should be registered in the buffer.</p>
<p>Thanks so much for your help!</p>

username: ptrblck
cooked: <p>You could call <code>.cuda()</code> or <code>to()</code>, but should specify the right device to push the tensor to.<br>
E.g. if you would like to create some tensors inside the <code>forward</code> method, you could use the device of some buffers/parameters or the incoming tensor to create the new one.</p>
<p>However, if <code>self.grid</code> is treated as an attribute of the model, registering it as a buffer is the cleaner and better approach. <img src="https://discuss.pytorch.org/images/emoji/apple/wink.png?v=9" title=":wink:" class="emoji" alt=":wink:"></p>

username: sunshineatnoon
cooked: <p>Right, that makes so much sense now. Thanks!</p>