How does Module.cuda() work and how to extend new modules that can use .cuda() to put its data to gpu?

username: Jiacheng_Li
cooked: <p>I want to write a loss function which will create new variables according to the shape of input data. The matter is when I use  my loss with <code>loss.cuda(0)</code>, the data in it won’t move to gpu device, and if the input data is on gpu it will occur error.</p>
<p>I referred to the code of loss functions already exists in PyTorch and noticed there is a class named _WeightedLoss and is used in NLLLoss. But because the weight NLL use is created in its <strong>init</strong> and needn’t adjust according to the input data, I don’t think this way is suitable for my problem.</p>
<p>The function I want to write can approximately described like below:</p>
<pre><code class="lang-auto">class exampleLoss(nn.Module):
    def __init__(self):
        super(exampleLoss,self).__init__()
    def forward(self,input, groundTruth):
        data1=torch.zeros_like(input)
        #some computation
        return result
</code></pre>
<p>I also noticed that the net I wrote before which also inherit from <code>nn.Module</code> can use <code>net.cuda()</code> to put all the layers and parameters to gpu even I don’t use functions like <code>nn.ModuleList</code> or <code>nn.register_buffer</code>. What’s the mechanism PyTorch used for collecting members of one class and make them available for <code>cuda()</code> operation?</p>
<p>Briefly, I want to use <code>.cuda()</code> to put my loss function which will create variables during <code>forward()</code> to gpu and is curious about the mechanism behind <code>Module.cuda()</code> operation.</p>
<p>Any help is appreciated. Thanks!</p>

username: albanD
cooked: <p>Hi,</p>
<p>When you call <code>.cuda()</code> all the parameters and buffers of the module are moved to the gpu.<br>
Parameters are everything that you saved as <code>self.foo = nn.Parameter(args)</code>.<br>
Buffers are any Tensor that you saved on self as <code>self.bar = torch.rand(10)</code>.<br>
All this is done  in the <a href="https://github.com/pytorch/pytorch/blob/f03296299b4f39ffed00d94c1adb12255b6d74c0/torch/nn/modules/module.py#L539-L581" rel="nofollow noopener"><code>__setattr__</code></a> functions of the Module if you want to check the implementation details.</p>
<p>If you create new tensors during the forward, it’s your responsability to create them with the right type. The set of functions <code>torch.*_like(input)</code> and <code>input.new_*(new_tensor_size)</code> can be used to easily create new tensors with the same type/device(/size) as other tensors.</p>

username: Jiacheng_Li
cooked: <p>Thanks for your reply. I haven’t noticed the device of <code>torch.*_like(input)</code> is also the same as that of input until you remind me.  I will check <code>__setattr__</code> for details. And that’s the second time you answered me. Thanks again!<img src="https://discuss.pytorch.org/images/emoji/apple/grin.png?v=6" title=":grin:" class="emoji" alt=":grin:"></p>

username: vb000
cooked: <p>Hi,</p>
<p>Does <code>model_cuda = model_cpu.cuda()</code> also copy the <code>model_cpu.state_dict()</code> to <code>model_cuda</code>?</p>
<p>Thanks,<br>
Bandhav</p>

username: albanD
cooked: <p>Hi,</p>
<p>The <code>.cuda()</code> of is inplace for Modules. So you can simply do <code>model.cuda()</code> to transform model from cpu to cuda.</p>

username: vb000
cooked: <p>Oh right! I just got to know that! Interesting that <code>tensor.cuda()</code> returns a copy of the tensor but <code>model.cuda()</code> is in-place.</p>