code: import torch
  import torch.nn as nn
  import torch.nn.functional as F
  
  
  class MyAwesomeModel(nn.Module):
      def __init__(self, n_classes):
          super(MyAwesomeModel, self).__init__()
  
          self.feature_extractor = nn.Sequential(
              nn.Conv2d(in_channels=1, out_channels=6, kernel_size=4, stride=1),
              nn.Tanh(),
              nn.AvgPool2d(kernel_size=2),
              nn.Conv2d(in_channels=6, out_channels=16, kernel_size=4, stride=1),
              nn.Tanh(),
              nn.AvgPool2d(kernel_size=2),
              nn.Conv2d(in_channels=16, out_channels=120, kernel_size=4, stride=1),
              nn.Tanh(),
          )
  
          self.classifier = nn.Sequential(
              nn.Linear(in_features=120, out_features=84),
              nn.Tanh(),
              nn.Linear(in_features=84, out_features=n_classes),
          )
  
      def forward(self, x, return_features=False):
          x = self.feature_extractor(x)
          x = torch.flatten(x, 1)
          logits = self.classifier(x)
          probs = F.log_softmax(logits, dim=1)
          if return_features:
              return x
          else:
              return probs
  
apis: torch.nn.functional.log_softmax torch.nn.Linear torch.flatten torch.nn.Tanh torch.nn.Conv2d torch.nn.AvgPool2d