code: import os
  import torch
  from torch import nn
  from torch.nn import functional as F
  from torchvision import datasets, transforms
  
  from src.models.base import BaseModel
  
  
  class MNIST(BaseModel):
  
      def _setup(self):
  
          self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
          self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
          self.conv2_drop = nn.Dropout2d()
          self.fc1 = nn.Linear(320, 50)
          self.fc2 = nn.Linear(50, 10)
  
      def forward(self, x):
  
          x = F.relu(F.max_pool2d(self.conv1(x), 2))
          x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))
          x = x.view(-1, 320)
          x = F.relu(self.fc1(x))
          x = F.dropout(x, training=self.training)
          x = self.fc2(x)
  
          return x
  
      def get_criterion(self):
  
          return nn.CrossEntropyLoss()
  
      def get_data_gen(self, batch_size, train=True):
  
          dpath_data = os.path.join(
              os.path.dirname(__file__),
              '..',
              '..',
              'data'
          )
  
          train_loader = torch.utils.data.DataLoader(
              datasets.MNIST(dpath_data, train=train, download=True,
                             transform=transforms.Compose([
                                 transforms.ToTensor(),
                                 transforms.Normalize((0.1307,), (0.3081,))
                             ])),
              batch_size=batch_size, shuffle=True)
  
          return iter(train_loader)
  
apis: torch.nn.Linear torch.nn.Dropout2d torch.nn.functional.dropout torch.nn.CrossEntropyLoss torch.nn.Conv2d