from torch import nn


class MyAwesomeModel(nn.Module):
    def __init__(self):
        super().__init__()

        self.cnn = nn.Sequential(nn.Conv2d(in_channels=1, out_channels=5, kernel_size=3),
                                nn.ReLU(),
                                nn.Conv2d(in_channels=5, out_channels=3, kernel_size=3, stride=2)
                                )
        self.fc = nn.Sequential(nn.Linear(432, 100),
                                nn.ReLU(),
                                nn.Linear(100,10),
                                nn.LogSoftmax(dim=1)
                                )

    def forward(self, x):
        x = self.cnn(x).view(x.size(0), -1)
        return self.fc(x)


[(8, 'torch.nn.Conv2d', 'nn.Conv2d', False, 'from torch import nn\n'), (9, 'torch.nn.ReLU', 'nn.ReLU', False, 'from torch import nn\n'), (10, 'torch.nn.Conv2d', 'nn.Conv2d', False, 'from torch import nn\n'), (12, 'torch.nn.Linear', 'nn.Linear', False, 'from torch import nn\n'), (13, 'torch.nn.ReLU', 'nn.ReLU', False, 'from torch import nn\n'), (14, 'torch.nn.Linear', 'nn.Linear', False, 'from torch import nn\n'), (15, 'torch.nn.LogSoftmax', 'nn.LogSoftmax', False, 'from torch import nn\n')]