text: import torch
  import torch.nn as nn
  
  
  class LeNet(nn.Module):
      def __init__(self):
          super(LeNet, self).__init__()
          self.conv = nn.Sequential(
              nn.Conv2d(1, 16, 5),
              nn.ReLU(),
              nn.MaxPool2d(2, 2),
              nn.Conv2d(16, 32, 5),
              nn.ReLU(),
              nn.MaxPool2d(2, 2),
              nn.Flatten()
          )
          self.fc = nn.Sequential(
              # image size changes to 800 after flatten (channels = 3)
              # 512 (channels = 1)
              nn.Linear(512, 120),
              nn.ReLU(),
              nn.Linear(120, 84),
              nn.ReLU(),
              nn.Linear(84, 10)
          )
      
      def forward(self, x):
          x = self.conv(x)
          x = self.fc(x)
          return x
  
  
  """
  model = LeNet()
  print(model)
  
  LeNet(
    (conv): Sequential(
      (0): Conv2d(1, 16, kernel_size=(5, 5), stride=(1, 1))
      (1): ReLU()
      (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
      (3): Conv2d(16, 32, kernel_size=(5, 5), stride=(1, 1))
      (4): ReLU()
      (5): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
      (6): Flatten(start_dim=1, end_dim=-1)
    )
    (fc): Sequential(
      (0): Linear(in_features=512, out_features=120, bias=True)
      (1): ReLU()
      (2): Linear(in_features=120, out_features=84, bias=True)
      (3): ReLU()
      (4): Linear(in_features=84, out_features=10, bias=True)
    )
  )
  """