Method __init__

sat/sgm/modules/diffusionmodules/model.py:140–141

(self, in_channels)

    """to match AttnBlock usage"""

    def __init__(self, in_channels):
        super().__init__(dim=in_channels, heads=1, dim_head=in_channels)


class AttnBlock(nn.Module):