Model.cuda() does not convert all variables to cuda

Hmrishav_Bandyopadhy: <p>Hi, so i am trying to write an architecture where i <em>have to</em> convert entire models to cuda using model.cuda(). However, some of the elements are variables initialised in the <strong>init</strong>() loop of nn.Module() class. How do i convert them to cuda ? For example,</p>
<pre><code class="lang-auto">class Net(nn.Module):
    def __init__(self):
        self.xyz=torch.tensor([1,2,3,4...])   
        # Convert this to cuda without using .cuda() on tensor xyz, but by using model.cuda()
        .
        .
    def forward(self,x):
        .
        .
</code></pre>
<p>Is it possible to do this?</p>

Valerio_Biscione: <p><code>torch.tensor([1,2,3]).cuda()</code> <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>
<p>EDIT: Wait, I just read your comment! At what point in your script you want to push those tensors in cuda? Why don’t you want to use .cuda()?</p>

Hmrishav_Bandyopadhy: <p>I want to do this:</p>
<pre><code class="lang-auto">model=Net()
model.cuda() ##Convert everything to cuda()
</code></pre>
<p>I dont want to use .cuda() as later on i will be migrating the code to a CPU only device where i would need to test on CPU.</p>
<p>EDIT: I will train on GPU and test on CPU, so the weights will need to be <code>.cpu()</code></p>

Valerio_Biscione: <p>I see, I also run stuff both on the CPU and GPU. My solution generally is something like</p>
<pre><code class="lang-auto">def make_cuda(fun):
    return fun.cuda() if torch.cuda.is_available() else fun

class Net(nn.Module):
    def __init__(self):
        self.xyz = make_cuda(torch.tensor([1,2,3,4...]))
        # Convert this to cuda without using .cuda() on tensor xyz, but by using model.cuda()
        .
        .
    def forward(self,x):
        .
        .

</code></pre>
<p>This will automatically push stuff to cuda only if you are ona machine with a GPU.</p>

Hmrishav_Bandyopadhy: <p>Wont work, as i would be packaging the model and converting it to tf-lite and stuff…I figured out something that would work though:</p>
<pre><code class="lang-auto">model=model.cpu()
model.xyz=model.xyz.cpu()
</code></pre>
<p>While saving the weights</p>

Valerio_Biscione: <p>This will push tensors to cuda when you call model.cuda():</p>
<pre><code class="lang-auto">
import torch
import torch.nn as nn
class Net(nn.Module):
    def __init__(self):
        super().__init__()
        self.xyz = torch.tensor([1,2,3,4])

    def forward(self,x):
        pass

    def cuda(self):
        super().cuda()
        for k, v in self.__dict__.items():
            if isinstance(v, torch.Tensor):
                v.cuda()
       

a = Net()
a.cuda()
</code></pre>
<p>edit: I mistakenly wrote self.cuda() instead of super().cuda. Fixed.</p>

ptrblck: <p>While this approach would work, the proper way to register tensors inside an <code>nn.Module</code> would be to either use <code>nn.Parameter</code> (if this tensor requires gradients and should be trained) or via <code>self.register_buffer</code>. Both approaches will make sure that this tensor will be pushed to the specified device (in <code>model.to(device)</code>) and will also be added to the <code>state_dict</code> (which would be important if you want to save and load this model).</p>
<p>CC <a class="mention" href="/u/hmrishav_bandyopadhy">@Hmrishav_Bandyopadhy</a></p>

Hmrishav_Bandyopadhy: <p>Can you add an example of <code>self.register_buffer</code> for this use-case ?</p>

Kitsunetic: <pre><code class="lang-auto">class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()

        xyz = torch.tensor([1, 2, 3, 4], dtype=torch.float)
        self.xyz = nn.Parameter(xyz)

    def forward(self, x):
        return x

model = Net()
model.xyz
"""
Parameter containing:
tensor([1., 2., 3., 4.], requires_grad=True)
"""

model.cuda()
model.xyz
"""
Parameter containing:
tensor([1., 2., 3., 4.], device='cuda:0', requires_grad=True)
"""
</code></pre>
<pre><code class="lang-auto">class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()

        xyz = torch.tensor([1, 2, 3, 4], dtype=torch.float)
        self.register_buffer("xyz", xyz)

    def forward(self, x):
        return x

model = Net()
model.xyz
"""
tensor([1., 2., 3., 4.])
"""

model.cuda()
model.xyz
"""
tensor([1., 2., 3., 4.], device='cuda:0')
"""

</code></pre>