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

[((169, 223), 'torch.nn.Conv2d', 'nn.Conv2d', (['(3)'], {'out_channels': '(6)', 'kernel_size': '(5)', 'padding': '(0)'}), '(3, out_channels=6, kernel_size=5, padding=0)\n', (178, 223), True, 'import torch.nn as nn\n'), ((258, 295), 'torch.nn.MaxPool2d', 'nn.MaxPool2d', ([], {'kernel_size': '(2)', 'stride': '(2)'}), '(kernel_size=2, stride=2)\n', (270, 295), True, 'import torch.nn as nn\n'), ((317, 372), 'torch.nn.Conv2d', 'nn.Conv2d', (['(6)'], {'out_channels': '(16)', 'kernel_size': '(5)', 'padding': '(0)'}), '(6, out_channels=16, kernel_size=5, padding=0)\n', (326, 372), True, 'import torch.nn as nn\n'), ((407, 444), 'torch.nn.MaxPool2d', 'nn.MaxPool2d', ([], {'kernel_size': '(2)', 'stride': '(2)'}), '(kernel_size=2, stride=2)\n', (419, 444), True, 'import torch.nn as nn\n'), ((464, 490), 'torch.nn.Linear', 'nn.Linear', (['(16 * 5 * 5)', '(120)'], {}), '(16 * 5 * 5, 120)\n', (473, 490), True, 'import torch.nn as nn\n'), ((506, 524), 'torch.nn.Linear', 'nn.Linear', (['(120)', '(84)'], {}), '(120, 84)\n', (515, 524), True, 'import torch.nn as nn\n'), ((544, 561), 'torch.nn.Linear', 'nn.Linear', (['(84)', '(10)'], {}), '(84, 10)\n', (553, 561), True, 'import torch.nn as nn\n'), ((693, 712), 'torch.flatten', 'torch.flatten', (['x', '(1)'], {}), '(x, 1)\n', (706, 712), False, 'import torch\n')]