Caching constant torch.eye at a nn.Module and moving to(device)

knownad: <p>I use a <code>torch.eye(.)</code> during the computation of a function. As it doesn’t change shape, so I decided to precompute it within my class once at the constructor instead of generating it on-the-fly during the <code>forward</code> call. Dummy example:</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn

class Foo(nn.Module):
    def __init__(self, device='cuda'):
        super(Foo, self).__init__()
        self.weights = nn.Parameter(torch.Tensor(4, 4))
        self.eye = torch.eye(4, device=device, requires_grad=False)
        torch.nn.init.uniform_(self.weights)

    def forward(self):
        return self.weights - self.eye
</code></pre>
<p>I have only added the weights as <code>nn.Parameter</code> so it won’t return <code>self.eye</code> when I call <code>model.parameters()</code> and added  <code>requires_grad=False</code> to make sure.</p>
<p>Should I be worried of any unintended behavior this could have in my computational graph?</p>
<p>This solution in contrast to this:</p>
<pre><code class="lang-auto">class Foo(nn.Module):
    def __init__(self, device='cuda'):
        super(Foo, self).__init__()
        self.weights = nn.Parameter(torch.Tensor(4, 4))
        self.device = device
        torch.nn.init.uniform_(self.weights)

    def forward(self):
        return self.weights - torch.eye(4, device=self.device)
</code></pre>
<p>Another question, when moving my model to a device, e.g.,</p>
<pre><code class="lang-auto">model = Foo()
model.to('cuda')
</code></pre>
<p>it correctly moves <code>self.weights</code> to the GPU, but not <code>self.eye</code> (PyTorch version 1.1.0). So I had to add a device parameter manually to ensure it was in the right device. Is it due to not being a <code>nn.Parameter</code>?</p>

albanD: <p>Hi,</p>
<p>This is the right way to do it. We call such Tensors buffers.<br>
It won’t have any impact on the computational graph as it does not require gradients.</p>
<p>It should be moved along with the parameters when you do <code>model.to('cuda')</code> though in your first code sample.</p>

knownad: <p>Thank you, <a class="mention" href="/u/alband">@albanD</a>!</p>
<p>About moving to GPU, I accidentally left the <code>device</code> argument. After removing it and trying to move <code>to(device)</code> again, it didn’t work in Pytorch 1.1.0:</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn

class Foo(nn.Module):
    def __init__(self):
        super(Foo, self).__init__()
        self.weights = nn.Parameter(torch.Tensor(4, 4))
        self.eye = torch.eye(4, requires_grad=False)
        torch.nn.init.uniform_(self.weights)

    def forward(self):
        return self.weights - self.eye

bar = Foo()
bar = bar.to('cuda')
print(bar.weights)  # device='cuda'
print(bar.eye.device)  # 'cpu'
bar()  # error
</code></pre>
<p>I also tried in Pytorch 1.3.0, and it also doesn’t move <code>self.eye</code> to GPU, only <code>self.weights</code>.</p>
<p>I’m currently keeping the <code>device</code> argument, just in case.</p>

albanD: <p>Ho my bad actually, you have to call it explicitly:<br>
<code>self.register_buffer('eye', torch.eye(4, requires_grad=False))</code> (instead of <code>self.eye = xxx</code>).</p>