import torch
import torch.nn as nn
import torch.nn.functional as F


class MediumNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(
            3, out_channels=6, kernel_size=5, padding=0)
        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.conv2 = nn.Conv2d(
            6, out_channels=16, kernel_size=5, padding=0)
        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.fc1 = nn.Linear(16*5*5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x = self.pool1(F.relu(self.conv1(x)))
        x = self.pool2(F.relu(self.conv2(x)))
        x = torch.flatten(x, 1)  # flatten all dimensions except batch
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

[((9, 21, 10, 56), 'torch.nn.Conv2d', 'nn.Conv2d', (), '', True, 'import torch.nn as nn\n'), ((11, 21, 11, 58), 'torch.nn.MaxPool2d', 'nn.MaxPool2d', (), '', True, 'import torch.nn as nn\n'), ((12, 21, 13, 57), 'torch.nn.Conv2d', 'nn.Conv2d', (), '', True, 'import torch.nn as nn\n'), ((14, 21, 14, 58), 'torch.nn.MaxPool2d', 'nn.MaxPool2d', (), '', True, 'import torch.nn as nn\n'), ((15, 19, 15, 41), 'torch.nn.Linear', 'nn.Linear', ({(15, 29, 15, 35): '16 * 5 * 5', (15, 37, 15, 40): '120'}, {}), '(16 * 5 * 5, 120)', True, 'import torch.nn as nn\n'), ((16, 19, 16, 37), 'torch.nn.Linear', 'nn.Linear', ({(16, 29, 16, 32): '120', (16, 34, 16, 36): '84'}, {}), '(120, 84)', True, 'import torch.nn as nn\n'), ((17, 19, 17, 36), 'torch.nn.Linear', 'nn.Linear', ({(17, 29, 17, 31): '84', (17, 33, 17, 35): '10'}, {}), '(84, 10)', True, 'import torch.nn as nn\n'), ((22, 12, 22, 31), 'torch.flatten', 'torch.flatten', ({(22, 26, 22, 27): 'x', (22, 29, 22, 30): '1'}, {}), '(x, 1)', False, 'import torch\n')]