class_docstr: ## LeNet based model fro MNIST classification
class_docstr_tok: ["#","#","LeNet","based","model","fro","MNIST","classification"]
docstr: * `x` is the batch of MNIST images of shape `[batch_size, 1, 28, 28]`
docstr_tok: ["*","`","x","`","is","the","batch","of","MNIST","images","of","shape","`","[","batch_size",",","1",",","28",",","28","]","`"]
code: def __call__(self, x: torch.Tensor):
      """
      * `x` is the batch of MNIST images of shape `[batch_size, 1, 28, 28]`
      """
      # Apply first convolution and max pooling.
      # The result has shape `[batch_size, 20, 12, 12]`
      x = self.max_pool1(self.act1(self.conv1(x)))
      # Apply second convolution and max pooling.
      # The result has shape `[batch_size, 50, 4, 4]`
      x = self.max_pool2(self.act2(self.conv2(x)))
      # Flatten the tensor to shape `[batch_size, 50 * 4 * 4]`
      x = x.view(x.shape[0], -1)
      # Apply hidden layer
      x = self.act3(self.fc1(x))
      # Apply dropout
      x = self.dropout(x)
      # Apply final layer and return
      return self.fc2(x)
  
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 2 for calling the Model obj's underlying member method code with example usage: obj.__call__(x) and returns: self