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


class ImageClassifierModel([nn.Module](https://docs.pytorch.org/docs/stable/generated/torch.nn.Module.html#torch.nn.Module "torch.nn.Module")):
    def __init__(self):
        super().__init__()
        self.conv1 = [nn.Conv2d](https://docs.pytorch.org/docs/stable/generated/torch.nn.Conv2d.html#torch.nn.Conv2d "torch.nn.Conv2d")(1, 6, 5)
        self.conv2 = [nn.Conv2d](https://docs.pytorch.org/docs/stable/generated/torch.nn.Conv2d.html#torch.nn.Conv2d "torch.nn.Conv2d")(6, 16, 5)
        self.fc1 = [nn.Linear](https://docs.pytorch.org/docs/stable/generated/torch.nn.Linear.html#torch.nn.Linear "torch.nn.Linear")(16 * 5 * 5, 120)
        self.fc2 = [nn.Linear](https://docs.pytorch.org/docs/stable/generated/torch.nn.Linear.html#torch.nn.Linear "torch.nn.Linear")(120, 84)
        self.fc3 = [nn.Linear](https://docs.pytorch.org/docs/stable/generated/torch.nn.Linear.html#torch.nn.Linear "torch.nn.Linear")(84, 10)

    def forward(self, x: [torch.Tensor](https://docs.pytorch.org/docs/stable/tensors.html#torch.Tensor "torch.Tensor")):
        x = [F.max_pool2d](https://docs.pytorch.org/docs/stable/generated/torch.nn.functional.max_pool2d.html#torch.nn.functional.max_pool2d "torch.nn.functional.max_pool2d")([F.relu](https://docs.pytorch.org/docs/stable/generated/torch.nn.functional.relu.html#torch.nn.functional.relu "torch.nn.functional.relu")(self.conv1(x)), (2, 2))
        x = [F.max_pool2d](https://docs.pytorch.org/docs/stable/generated/torch.nn.functional.max_pool2d.html#torch.nn.functional.max_pool2d "torch.nn.functional.max_pool2d")([F.relu](https://docs.pytorch.org/docs/stable/generated/torch.nn.functional.relu.html#torch.nn.functional.relu "torch.nn.functional.relu")(self.conv2(x)), 2)
        x = [torch.flatten](https://docs.pytorch.org/docs/stable/generated/torch.flatten.html#torch.flatten "torch.flatten")(x, 1)
        x = [F.relu](https://docs.pytorch.org/docs/stable/generated/torch.nn.functional.relu.html#torch.nn.functional.relu "torch.nn.functional.relu")(self.fc1(x))
        x = [F.relu](https://docs.pytorch.org/docs/stable/generated/torch.nn.functional.relu.html#torch.nn.functional.relu "torch.nn.functional.relu")(self.fc2(x))
        x = self.fc3(x)
        return x