# LeNet Model definition
class Net([nn.Module](https://docs.pytorch.org/docs/stable/generated/torch.nn.Module.html#torch.nn.Module "torch.nn.Module")):
    def __init__(self):
        super([Net](https://docs.pytorch.org/docs/stable/generated/torch.nn.Module.html#torch.nn.Module "torch.nn.Module"), self).__init__()
        self.conv1 = [nn.Conv2d](https://docs.pytorch.org/docs/stable/generated/torch.nn.Conv2d.html#torch.nn.Conv2d "torch.nn.Conv2d")(1, 32, 3, 1)
        self.conv2 = [nn.Conv2d](https://docs.pytorch.org/docs/stable/generated/torch.nn.Conv2d.html#torch.nn.Conv2d "torch.nn.Conv2d")(32, 64, 3, 1)
        self.dropout1 = [nn.Dropout](https://docs.pytorch.org/docs/stable/generated/torch.nn.Dropout.html#torch.nn.Dropout "torch.nn.Dropout")(0.25)
        self.dropout2 = [nn.Dropout](https://docs.pytorch.org/docs/stable/generated/torch.nn.Dropout.html#torch.nn.Dropout "torch.nn.Dropout")(0.5)
        self.fc1 = [nn.Linear](https://docs.pytorch.org/docs/stable/generated/torch.nn.Linear.html#torch.nn.Linear "torch.nn.Linear")(9216, 128)
        self.fc2 = [nn.Linear](https://docs.pytorch.org/docs/stable/generated/torch.nn.Linear.html#torch.nn.Linear "torch.nn.Linear")(128, 10)

    def forward(self, x):
        x = self.conv1(x)
        x = [F.relu](https://docs.pytorch.org/docs/stable/generated/torch.nn.functional.relu.html#torch.nn.functional.relu "torch.nn.functional.relu")(x)
        x = self.conv2(x)
        x = [F.relu](https://docs.pytorch.org/docs/stable/generated/torch.nn.functional.relu.html#torch.nn.functional.relu "torch.nn.functional.relu")(x)
        x = [F.max_pool2d](https://docs.pytorch.org/docs/stable/generated/torch.nn.functional.max_pool2d.html#torch.nn.functional.max_pool2d "torch.nn.functional.max_pool2d")(x, 2)
        x = self.dropout1(x)
        x = [torch.flatten](https://docs.pytorch.org/docs/stable/generated/torch.flatten.html#torch.flatten "torch.flatten")(x, 1)
        x = self.fc1(x)
        x = [F.relu](https://docs.pytorch.org/docs/stable/generated/torch.nn.functional.relu.html#torch.nn.functional.relu "torch.nn.functional.relu")(x)
        x = self.dropout2(x)
        x = self.fc2(x)
        output = [F.log_softmax](https://docs.pytorch.org/docs/stable/generated/torch.nn.functional.log_softmax.html#torch.nn.functional.log_softmax "torch.nn.functional.log_softmax")(x, dim=1)
        return output

# MNIST Test dataset and dataloader declaration
[test_loader](https://docs.pytorch.org/docs/stable/data.html#torch.utils.data.DataLoader "torch.utils.data.DataLoader") = [torch.utils.data.DataLoader](https://docs.pytorch.org/docs/stable/data.html#torch.utils.data.DataLoader "torch.utils.data.DataLoader")(
    [datasets.MNIST](https://docs.pytorch.org/vision/stable/generated/torchvision.datasets.MNIST.html#torchvision.datasets.MNIST "torchvision.datasets.MNIST")('../data', train=False, download=True, transform=[transforms.Compose](https://docs.pytorch.org/vision/stable/generated/torchvision.transforms.Compose.html#torchvision.transforms.Compose "torchvision.transforms.Compose")([
            [transforms.ToTensor](https://docs.pytorch.org/vision/stable/generated/torchvision.transforms.ToTensor.html#torchvision.transforms.ToTensor "torchvision.transforms.ToTensor")(),
            [transforms.Normalize](https://docs.pytorch.org/vision/stable/generated/torchvision.transforms.Normalize.html#torchvision.transforms.Normalize "torchvision.transforms.Normalize")((0.1307,), (0.3081,)),
            ])),
        batch_size=1, shuffle=True)

# We want to be able to train our model on an `accelerator <https://pytorch.org/docs/stable/torch.html#accelerators>`__
# such as CUDA, MPS, MTIA, or XPU. If the current accelerator is available, we will use it. Otherwise, we use the CPU.
device = [torch.accelerator.current_accelerator](https://docs.pytorch.org/docs/stable/generated/torch.accelerator.current_accelerator.html#torch.accelerator.current_accelerator "torch.accelerator.current_accelerator")().type if [torch.accelerator.is_available](https://docs.pytorch.org/docs/stable/generated/torch.accelerator.is_available.html#torch.accelerator.is_available "torch.accelerator.is_available")() else "cpu"
print(f"Using {device} device")

# Initialize the network
model = [Net](https://docs.pytorch.org/docs/stable/generated/torch.nn.Module.html#torch.nn.Module "torch.nn.Module")().to(device)

# Load the pretrained model
[model.load_state_dict](https://docs.pytorch.org/docs/stable/generated/torch.nn.Module.html#torch.nn.Module.load_state_dict "torch.nn.Module.load_state_dict")([torch.load](https://docs.pytorch.org/docs/stable/generated/torch.load.html#torch.load "torch.load")(pretrained_model, map_location=device, weights_only=True))

# Set the model in evaluation mode. In this case this is for the Dropout layers
[model.eval](https://docs.pytorch.org/docs/stable/generated/torch.nn.Module.html#torch.nn.Module.eval "torch.nn.Module.eval")()