def __call__()

in jcm/models/layers.py [0:0]

• 33 lines of code
• 7 McCabe index (conditional complexity)

    def __call__(self, x):
        h = self.normalization()(x)
        h = self.act(h)
        if self.resample == "down":
            h = ncsn_conv3x3(h, h.shape[-1], dilation=self.dilation)
            h = self.normalization()(h)
            h = self.act(h)
            if self.dilation > 1:
                h = ncsn_conv3x3(h, self.output_dim, dilation=self.dilation)
                shortcut = ncsn_conv3x3(x, self.output_dim, dilation=self.dilation)
            else:
                h = ConvMeanPool(output_dim=self.output_dim)(h)
                shortcut = ConvMeanPool(output_dim=self.output_dim, kernel_size=1)(x)
        elif self.resample is None:
            if self.dilation > 1:
                if self.output_dim == x.shape[-1]:
                    shortcut = x
                else:
                    shortcut = ncsn_conv3x3(x, self.output_dim, dilation=self.dilation)
                h = ncsn_conv3x3(h, self.output_dim, dilation=self.dilation)
                h = self.normalization()(h)
                h = self.act(h)
                h = ncsn_conv3x3(h, self.output_dim, dilation=self.dilation)
            else:
                if self.output_dim == x.shape[-1]:
                    shortcut = x
                else:
                    shortcut = ncsn_conv1x1(x, self.output_dim)
                h = ncsn_conv3x3(h, self.output_dim)
                h = self.normalization()(h)
                h = self.act(h)
                h = ncsn_conv3x3(h, self.output_dim)

        return h + shortcut