BentouAI

class SmallBlock(torch.nn.Module):
    def __init__(self, in_channels, out_channels):
        super(SmallBlock, self).__init__()
        self.model = torch.nn.Sequential(
                torch.nn.Conv2d(in_channels, out_channels, 3, stride=1, padding=1),
                torch.nn.BatchNorm2d(out_channels)
        )
    def forward(self, X):
        return self.model(X)