class_docstr: ## LeNet based model fro MNIST classification
class_docstr_tok: ["#","#","LeNet","based","model","fro","MNIST","classification"]
code: def __init__(self, dropout: float):
      super().__init__()
      # First $5x5$ convolution layer
      self.conv1 = nn.Conv2d(1, 20, kernel_size=5)
      # ReLU activation
      self.act1 = nn.ReLU()
      # $2x2$ max-pooling
      self.max_pool1 = nn.MaxPool2d(2, 2)
      # Second $5x5$ convolution layer
      self.conv2 = nn.Conv2d(20, 50, kernel_size=5)
      # ReLU activation
      self.act2 = nn.ReLU()
      # $2x2$ max-pooling
      self.max_pool2 = nn.MaxPool2d(2, 2)
      # First fully-connected layer that maps to $500$ features
      self.fc1 = nn.Linear(50 * 4 * 4, 500)
      # ReLU activation
      self.act3 = nn.ReLU()
      # Final fully connected layer to output evidence for $10$ classes.
      # The ReLU or Softplus activation is applied to this outside the model to get the
      # non-negative evidence
      self.fc2 = nn.Linear(500, 10)
      # Dropout for the hidden layer
      self.dropout = nn.Dropout(p=dropout)
  
imports: from typing import Any
  import torch.nn
  import torch.utils.data
  from labml import tracker, experiment
  from labml.configs import option, calculate
  from labml_helpers.module import Module
  from labml_helpers.schedule import Schedule, RelativePiecewise
  from labml_helpers.train_valid import BatchIndex
  from labml_nn.experiments.mnist import MNISTConfigs
  from labml_nn.uncertainty.evidence import KLDivergenceLoss, TrackStatistics, MaximumLikelihoodLoss, CrossEntropyBayesRisk, SquaredErrorBayesRisk
label_desc: Use image node_id 1 to create a new Model object from inherited base classes: Module with example: obj = Model(dropout)