Creating new layer

username: Mohamed_Farag
cooked: <p>Hi,<br>
I am trying to re-implement a channel attention module in order to get my hands-on dealing with pytorch and how to create my custom layers.</p>
<pre><code class="lang-auto">class CAM_infer(nn.Module):
    """Channel Attention module"""
    
    def __init__(self, ratio= 16):
        super(CAM_infer, self).__init__()
        # self.channels = channels
        self.ratio =  ratio

        self.avgpool = nn.AdaptiveAvgPool2d((1,1))
        self.maxpool = nn.AdaptiveMaxPool2d((1,1))

        # self.h1 = nn.LazyLinear(channels//self.ratio, bias=True)
        # self.h2 = nn.LazyLinear(channels, bias= True)

        # nn.init.xavier_uniform_(self.h1.weight)
        # nn.init.xavier_uniform_(self.h2.weight)

        self.act = nn.Sigmoid()

        

    def forward(self, x):

        no_channels = x.size(1)
        h1 = nn.LazyLinear(no_channels//self.ratio, bias=True)
        h2 = nn.LazyLinear(no_channels, bias= True)

        x1 = self.avgpool(x)
        x2 = self.maxpool(x)

        x1 = x1.view(-1, 1 ,no_channels)
        x2 = x2.view(-1, 1 ,no_channels)

        x1 = nn.functional.relu(h1(x1))
        x1 = h2(x1)

        x2 = nn.functional.relu(h1(x2))
        x2 = h2(x2)

        x_new = torch.add(x1, x2)
        mc = self.act(x_new)
        mc = mc.view(-1, no_channels * 1)

        # x = torch.permute(x, (2,3, 0, 1))

        # f_dash = torch.multiply(x, mc)

        # f_dash.permute((2, 3, 0, 1))
        
        return  mc
</code></pre>
<p>I want to infer the input for this layer from the previous layer, it is like a LazyLinear approach,</p>
<p>I run this test and it is totally running:</p>
<pre><code class="lang-auto">input = torch.randn(32, 1664, 7, 7)

# GAP = nn.AdaptiveAvgPool2d((1,1))

# GAP(input).shape

# torch.squeeze(input, 1).shape

cam = CAM_infer()

mc = cam(input)

if mc.ndim != 2:

raise ValueError('Output dimension should be 2-D')

print(mc.shape)
</code></pre>
<p>But when I transfer the data and the model to CUDA side I got this error:</p>
<p><code>RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cpu and cuda:0! (when checking argument for argument mat1 in method wrapper_CUDA_addmm)</code></p>
<p>May anyone help? <img src="https://discuss.pytorch.org/images/emoji/apple/eyes.png?v=12" title=":eyes:" class="emoji" alt=":eyes:" loading="lazy" width="20" height="20"></p>
<p>Thanks a lot for your continuous support <img src="https://discuss.pytorch.org/images/emoji/apple/heart.png?v=12" title=":heart:" class="emoji" alt=":heart:" loading="lazy" width="20" height="20"></p>

username: Mohamed_Farag
cooked: <p>it seems there is a problem regarding lazyLinear parts, cause in the forward method now I transfered them to the cuda side before creating an instance of the CAM_infer and it seems it is working properly <img src="https://discuss.pytorch.org/images/emoji/apple/thinking.png?v=12" title=":thinking:" class="emoji" alt=":thinking:" loading="lazy" width="20" height="20"></p>
<p>May any one elaborate?</p>

username: eqy
cooked: <p>Yes, the problem is expected to be with the <code>LazyLinear</code> layers because you are instantiating them repeatedly each time the forward pass executes, and by default they would have device CPU.</p>
<p>Your other layers are correctly moved to CUDA because they were declared in the <code>__init__</code> method, and calling <code>.cuda()</code> on the module should move these layers. But as your <code>LazyLinear</code> layers were not already created, they would not be moved.</p>
<p>Note that this also causes another problem—training this model wouldn’t actually train the <code>LazyLinear</code> layers because they are being recreated and reinitialized every forward pass.</p>