Module.cuda() not moving Module tensor?

username: danh
cooked: <pre><code class="lang-auto">class ToyModule(torch.nn.Module):
    def __init__(self) -&gt; None:
        super(ToyModule, self).__init__()
        self.layer = torch.nn.Linear(2, 2)
        self.expected_moved_cuda_tensor = torch.tensor([0, 2, 3])

    def forward(self, input: torch.Tensor) -&gt; torch.Tensor:
        return self.layer(input)
</code></pre>
<pre><code class="lang-auto">toy_module = ToyModule()
toy_module.cuda()
</code></pre>
<p>When we call  <code>.cuda()</code>  all the parameters and buffers of the module are moved to the GPU:</p>
<pre><code class="lang-auto">next(toy_module.layer.parameters()).device
&gt;&gt;&gt; device(type='cuda', index=0)
</code></pre>
<p>But when we inspect the tensor attribute of toy_module, we see <code>device(type='cpu')</code>?</p>
<pre><code class="lang-auto">toy_module.expected_moved_cuda_tensor.device
&gt;&gt;&gt; device(type='cpu')
</code></pre>
<p>Is this expected or am I missing anything? Thank you.</p>

username: braindotai
cooked: <aside class="quote no-group" data-username="danh" data-post="1" data-topic="74582">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/d/a9a28c/48.png" class="avatar"> danh:</div>
<blockquote>
<p>When we call <code>.cuda()</code> all the parameters and buffers of the module are moved to the GPU:</p>
</blockquote>
</aside>
<p><code>self.expected_moved_cuda_tensor</code> is neither a parameter nor a buffer, that’s why it’s device is unchanged. If you want to create a parameter and use it then you can do it as follows-</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
import torch.nn as nn

class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.linear1 = nn.Linear(2, 1)
        self.linear1.weight = torch.nn.Parameter(torch.ones(2, 1))
        self.linear1.bias = torch.nn.Parameter(torch.zeros(1))

    def forward(self, x):
        x = self.linear1(x)
        return x
</code></pre>
<p>You can even use those parameters in forward method like-</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
import torch.nn as nn

class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.weight = torch.nn.Parameter(torch.ones(2, 1))
        self.bias = torch.nn.Parameter(torch.zeros(1))

    def forward(self, x):
        # linear regression completely from scratch,
        # using parameters created in __init__
        x = torch.mm(x, self.weight) + self.bias
        return x
</code></pre>
<p>And moving above model to <code>.cuda()</code> does move model parameters to cuda-</p>
<pre data-code-wrap="python"><code class="lang-python">model = Model()
model.cuda()
print(model.weight.device) # prints device(type='cuda', index=0)
print(model.bias.device) # prints device(type='cuda', index=0)</code></pre>

username: danh
cooked: <p>Thanks a lot!</p>
<p>But isn’t it defeat the intuition of  <code>.cuda()</code>  if the Module tensor device stays unchanged?</p>

username: braindotai
cooked: <p>Though <code>.cuda()</code> “should” do as you said, but I don’t think changing devices for all the <code>torch.tensor </code> attributes of a class inherited from <code>nn.Module</code> by default is good idea. In your use case it might be helpful, but in some case user may don’t that, so I think that’s the reason why it ain’t do that by default.</p>
<p>One more thing if you want to create just a constant tensor (not a parameter) then you can do that as</p>
<pre><code class="lang-python">self.a_constant_tensor = nn.Parameter(torch.ones(2, 1), requires_grad = False)
</code></pre>
<p>and then use it in forward method.</p>
<p><strong>OR</strong> you can use buffers, <strong>"which is recommended"</strong></p>
<pre><code class="lang-python">import torch
import torch.nn as nn

class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.weight = torch.nn.Parameter(torch.zeros(2, 1))
        self.bias = torch.nn.Parameter(torch.zeros(1))
        self.register_buffer('a_constant_tensor', torch.tensor([0.5]))

    def forward(self, x):
        # linear regression completely from scratch,
        # using parameters created in __init__
        x = torch.mm(x, self.weight) + self.bias + self.a_constant_tensor
        return x


model = Model().cuda()
</code></pre>
<p>Doing this wouldn’t consider <code>self.a_constant_tensor</code> as a parameter, so printing parameters wouldn’t return <code>self.a_constant_tensor</code> -</p>
<pre><code class="lang-python">for param in model.parameters():
    print(param)
# Only prints about self.weight and self.bias
'''
Parameter containing:
tensor([[0.],
        [0.]], device='cuda:0', requires_grad=True)
Parameter containing:
tensor([0.], device='cuda:0', requires_grad=True)
'''
</code></pre>