Method __init__

models/unet_model_google.py:97–106

(self, dim, dim_out, groups=32, dropout=0)

class Block(nn.Module):
    def __init__(self, dim, dim_out, groups=32, dropout=0):
        super().__init__()
        self.block = nn.Sequential(
            # nn.BatchNorm2d(dim),
            nn.GroupNorm(groups, dim),
            # LayerNorm2d(dim),
            Swish(),
            nn.Dropout(dropout) if dropout != 0 else nn.Identity(),
            nn.Conv2d(dim, dim_out, 3, padding=1),
        )

    def forward(self, x):
        return self.block(x)