class_docstr: MNIST classifier used in the experiments for Counterfactual with Reinforcement Learning. The model consists of two
  convolutional layers having 64 and 32 channels and a kernel size of 2 with ReLU nonlinearities, followed by
  maxpooling of size 2 and dropout of 0.3. The convolutional block is followed by a fully connected layer of 256 with
  ReLU nonlinearity, and finally a fully connected layer is used to predict the class logits (10 in MNIST case).
docstr: Constructor.
  
  Parameters
  ----------
  output_dim
      Output dimension.
docstr_tok: ["Constructor",".","Parameters","--","--","--","--","--","output_dim","Output","dimension","."]
code: def __init__(self, output_dim: int) -> None:
      """
      Constructor.
  
      Parameters
      ----------
      output_dim
          Output dimension.
      """
      super().__init__()
  
      self.conv1 = nn.Conv2d(1, 64, kernel_size=(2, 2), padding=1)
      self.maxpool1 = nn.MaxPool2d(kernel_size=2, stride=2)
      self.dropout1 = nn.Dropout(0.3)
      self.conv2 = nn.Conv2d(64, 32, kernel_size=(2, 2), padding=1)
      self.maxpool2 = nn.MaxPool2d(kernel_size=2, stride=2)
      self.dropout2 = nn.Dropout(0.3)
      self.flatten = nn.Flatten()
      self.fc1 = nn.Linear(1568, 256)
      self.fc2 = nn.Linear(256, output_dim)
  
      # send to device
      self.to(self.device)
  
imports: import torch
  import torch.nn
  import torch.nn.functional
  from typing import List
  from alibi.models.pytorch.model import Model
label_desc: Use image node_id 1 to create a new MNISTClassifier object from inherited base classes: Model with example: obj = MNISTClassifier(output_dim)