CycleFC implementation crashes during training

username: Ahmad_Omar_Ahsan
cooked: <p>Hello everyone,<br>
I am trying to implement CycleFC module from this paper <a href="https://arxiv.org/abs/2107.10224" rel="noopener nofollow ugc">CycleMLP</a>. Here is my code.</p>
<pre><code class="lang-auto">class CycleFC(nn.Module):
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        sh: int,
        sw: int,
    ):
        """Initialize CycleFC layer

        Args:
            in_channels (int): Input_channels.
            out_channels (int): Ouput_channels.
            sh (int): Stepsize along height.
            sw (int): Stepsize along width.

        """
        super(CycleFC, self).__init__()
        self.W_mlp = nn.parameter.Parameter(
            torch.randn(size=(in_channels, out_channels)),
            requires_grad=True,
        )
        self.bias = nn.parameter.Parameter(torch.randn(size=(out_channels,)))

        self.sh = sh
        self.sw = sw
        self.in_channels = in_channels
        self.out_channels = out_channels

    def forward(self, x: torch.Tensor) -&gt; torch.Tensor:
        """Forward function for CycleFC

        Args:
            x (torch.Tensor): Input Tensor. Shape : b,  c_in, h, w

        Returns:
            torch.Tensor: Output Tensor. Shape : b, c_out, h, w
        """
        b, c_in, h, w = x.shape
        output = torch.zeros(size=(b, self.out_channels, h, w)).to(x.device)
        for k in range(b):
            for i in range(h):
                for j in range(w):
                    output[k, :, i, j] = self.calc_value(
                        x, w_mlp=self.W_mlp, c_in=c_in, i=i, j=j, h=h, w=w, k=k
                    )
                    output[k, :, i, j] += self.bias
        return output

    def get_offset(self, sh: int, sw: int, c: int) -&gt; Tuple[int, int]:
        """Calculate offset based on stepsize along height and width for a given channel.

        Args:
            sh (int): Stepsize along height
            sw (int): Stepsize along width
            c (int): Current channel

        Returns:
            Tuple[int, int]: Returns offset along height and width for a given channel.
        """
        delta_i = c % sh - 1
        delta_j = floor(c / sh) % sw - 1
        return delta_i, delta_j

    def calc_value(
        self,
        x: torch.Tensor,
        w_mlp: torch.Tensor,
        c_in: int,
        i: int,
        j: int,
        h: int,
        w: int,
        k: int,
    ):
        """Calculate value for CycleFC

        Args:
            x (torch.Tensor): Input Tensor. Shape : b, c, h, w,
            w_mlp (torch.Tensor): Weight matrix.
            c_in (int): Input Channel
            i (int) : Current Height index
            j(int) : Current width index
            h(int) : Input height
            w(int) : Input weight
            k(int) : Batch size
        Returns:
            torch.Tensor: Output tensor
        """
        sum = 0
        for c in range(c_in):
            delta_i, delta_j = self.get_offset(self.sh, self.sw, c)
            i_offset = (i + delta_i) % h
            j_offset = (j + delta_j) % w
            sum += x[k, c, i_offset, j_offset] * w_mlp[c, :]
        return sum

</code></pre>
<p>This is a naive implementation of the formula provided in the paper for CycleFC. My training crashes, and it is extremely slow. How can I make this better? I have also attached a screenshot of the formula.<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/6/1/61cd6e72796691952de1bc66cd33fd6db00e351e.png" data-download-href="https://discuss.pytorch.org/uploads/default/61cd6e72796691952de1bc66cd33fd6db00e351e" title="cyclefc"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/6/1/61cd6e72796691952de1bc66cd33fd6db00e351e_2_690x429.png" alt="cyclefc" data-base62-sha1="dXcrRGSRO9Luraqbzr5UlFs4xR4" width="690" height="429" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/6/1/61cd6e72796691952de1bc66cd33fd6db00e351e_2_690x429.png, https://discuss.pytorch.org/uploads/default/original/3X/6/1/61cd6e72796691952de1bc66cd33fd6db00e351e.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/6/1/61cd6e72796691952de1bc66cd33fd6db00e351e.png 2x" data-dominant-color="EDEDED"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">cyclefc</span><span class="informations">931×579 94.8 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>