def __init__()

in codes/models.py [0:0]

• 8 lines of code
• 1 McCabe index (conditional complexity)

    def __init__(self, dim, pool_size = 8):
        super().__init__()
        self.pool_size = pool_size
        self.pool_fn = Rearrange('b d (n p) -> b d n p', p = pool_size)

        self.to_attn_logits = nn.Conv2d(dim, dim, 1, bias = False)

        nn.init.dirac_(self.to_attn_logits.weight)

        with torch.no_grad():
            self.to_attn_logits.weight.mul_(2)