Confusion about torch.cuda.device context mamanger

kirk86: <p>Hey folks,<br>
I was wondering if someone could clarify some aspects of the cuda context manager.</p>
<p>For instance lets say we have dummy function, and we want to either send the whole function to cuda or whatever tensors are created/manipulated inside that dummy function to automatically get sent to cuda.</p>
<pre><code class="lang-auto">def dummy_func(input, output):
    y = input * output
    tmp = torch.ones(input.size(0), output.size(2))
    y.scatter_(1, tmp, 1)
</code></pre>
<p>What I would like is everythin inside <code>dummy_func</code> to operate on cuda without having to explicity add to each new tensor <code>.to(device)</code> or <code>.cuda()</code></p>
<p>What I’ve tried are the following</p>
<pre><code class="lang-auto">@torch.cuda()
def dummy_func(...):
</code></pre>
<p>and</p>
<pre><code class="lang-auto">def dummy_func(....):
    with torch.cuda.device(0):
           .....
</code></pre>
<p><img src="https://discuss.pytorch.org/uploads/default/original/3X/2/b/2b5c4b7695fc28aa3f37d734e765855627deb14c.png" alt="image" data-base62-sha1="6bAgR6xEFKPujNkdwBU4oU2SZEU" width="408" height="113"></p>
<p>But none of those gave the expected outcome.<br>
Is there any way to achieve that in order to avoid having to manually add <code>.cuda</code> to every possible tensor?</p>

albanD: <p>Hi,</p>
<p>The <code>torch.cuda.device()</code> context is to change the defaut cuda device so that any cuda-related function will use that device.</p>
<p>We don’t have anything that forces all Tensors to move to cuda. This is mainly because this is a quite expensive operation and the user should be aware of it and avoid moving Tensors back and forth between the CPU and GPU.</p>

kirk86: <p>I understand but think of the following scenario.<br>
We have a model named <code>net(inputs, targets)</code> taking inputs and targets.<br>
So far so good things are easy we can send <code>net.to(device)</code> and <code>inputs.to(device), targets.to(device)</code> accordingly.</p>
<p>Things start to get hairier when we want to manipulate some aspects of model as shown in the example above in <code>def dummy_func</code>.</p>
<p>Every <code>torch.ones, torch.zeros</code>  whatnot has to be accompanied by <code>.cuda</code>.</p>
<p>Which most of the time will raise errors if we by mistake have at least one tensor that is on cpu and the rest on gpu and vice versa.</p>
<p>I think this defeats the purpose of <code>model.to(device)</code>, does’t it?</p>
<p>We would expect that whatever is happening inside the model (whatever ops) to be sent to device, no?</p>

albanD: <p>But you will see the same problem if you change the input type.<br>
In general, you want to do:</p>
<pre><code class="lang-auto">tmp = torch.ones(input.size(0), output.size(2), dtype=input.dtype, device=input.device)
</code></pre>
<p>to make sure you get a Tensor of the same type and same device as input.<br>
If you have the same size (not the case for this example), you can do <code>tmp = torch.ones_like(input)</code> to get the same size/dtype/device.</p>

kirk86: <p>Ok, I didn’t know that functions <code>torch.ones_like</code> would automatically put things on same device. Thanks!<br>
But I think that we’re still missing something</p>
<pre><code class="lang-auto">with torch.cuda(device):
   whatever happens here make them automatically on cuda device defined
</code></pre>
<p>or something like</p>
<pre><code class="lang-auto">@torch.cuda(device)
def some_func():
  send whatever is in here to cuda device
</code></pre>
<p>Unless there exists something for such cases and I’ve completely missed it?</p>

albanD: <p>There isn’t but that’s by choice <img src="https://discuss.pytorch.org/images/emoji/apple/smiley.png?v=9" title=":smiley:" class="emoji" alt=":smiley:"></p>
<p>You could do something like <code>torch.set_default_tensor_type(torch.cuda.FloatTensor)</code> but this is <strong>strongly</strong> advised against. Anything you will create will be on the GPU (temporary stuff for printing, internal buffers,…) and you most likely don’t want these on the GPU as ops on them will be much slower.</p>