class_docstr: ## Model for classifying MNIST digits
class_docstr_tok: ["#","#","Model","for","classifying","MNIST","digits"]
docstr: `data` are the MNIST images, with shape `[batch_size, 1, 28, 28]`
docstr_tok: ["`","data","`","are","the","MNIST","images",",","with","shape","`","[","batch_size",",","1",",","28",",","28","]","`"]
code: def forward(self, data: torch.Tensor):
      """
      `data` are the MNIST images, with shape `[batch_size, 1, 28, 28]`
      """
      # Pass through the first convolution layer.
      # Output of this layer has shape `[batch_size, 256, 20, 20]`
      x = F.relu(self.conv1(data))
      # Pass through the second convolution layer.
      # Output of this has shape `[batch_size, 32 * 8, 6, 6]`.
      # *Note that this layer has a stride length of $2$*.
      x = self.conv2(x)
  
      # Resize and permutate to get the capsules
      caps = x.view(x.shape[0], 8, 32 * 6 * 6).permute(0, 2, 1)
      # Squash the capsules
      caps = self.squash(caps)
      # Take them through the router to get digit capsules.
      # This has shape `[batch_size, 10, 16]`.
      caps = self.digit_capsules(caps)
  
      # Get masks for reconstructioon
      with torch.no_grad():
          # The prediction by the capsule network is the capsule with longest length
          pred = (caps**2).sum(-1).argmax(-1)
          # Create a mask to maskout all the other capsules
          mask = torch.eye(10, device=data.device)[pred]
  
      # Mask the digit capsules to get only the capsule that made the prediction and
      # take it through decoder to get reconstruction
      reconstructions = self.decoder(
          (caps * mask[:, :, None]).view(x.shape[0], -1)
      )
      # Reshape the reconstruction to match the image dimensions
      reconstructions = reconstructions.view(-1, 1, 28, 28)
  
      return caps, reconstructions, pred
  
imports: from typing import Any
  import torch.nn
  import torch.nn.functional
  import torch.utils.data
  from labml import experiment, tracker
  from labml.configs import option
  from labml_helpers.datasets.mnist import MNISTConfigs
  from labml_helpers.metrics.accuracy import AccuracyDirect
  from labml_helpers.module import Module
  from labml_helpers.train_valid import SimpleTrainValidConfigs, BatchIndex
  from labml_nn.capsule_networks import Squash, Router, MarginLoss
label_desc: Use image node_id 2 for calling the MNISTCapsuleNetworkModel obj's underlying member method code with example usage: obj.forward(data) and returns: caps,  reconstructions,  pred