# Load the required libraries
import torch
import torch.nn as nn
from torch.utils.data import DataLoader, Dataset
from torchvision import datasets, transforms
import matplotlib.pyplot as plt
from PIL import Image
import os

# Device configuration
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# Define the CNN model exactly as in chapter 2.8
class CNN(nn.Module):

    def __init__(self):
        super(CNN, self).__init__()

        # Define model layers
        self.model_layers = nn.Sequential(

            nn.Conv2d(in_channels=1, out_channels=6, kernel_size=5),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2),

            nn.Conv2d(in_channels=6, out_channels=16, kernel_size=5),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2),

            nn.Flatten(),
            nn.Linear(16*97*172, 120),
            nn.ReLU(),
            nn.Linear(120, 2),
            #nn.LogSoftmax(dim=1)
        )

    def forward(self, x):
        out = self.model_layers(x)
        return out

# Load the model's parameters
model = torch.load("model.pth")
model.eval()