main-noncomp.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Prelims"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<contextlib.ExitStack at 0x297c72530>"
      ]
     },
     "execution_count": 47,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import os\n",
    "import torch\n",
    "import pandas as pd\n",
    "from skimage import io, transform\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "from torchvision import transforms, utils\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "\n",
    "# Ignore warnings\n",
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\")\n",
    "\n",
    "import logging\n",
    "\n",
    "debug = logging.getLogger(\"Debug\")\n",
    "info  = print\n",
    "plt.ion()   # interactive mode"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Utils"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Running MPS Mode: mps\n"
     ]
    }
   ],
   "source": [
    "#check GPU\n",
    "device = None\n",
    "if torch.cuda.is_available():\n",
    "    device = torch.device(\"cuda\")\n",
    "    print(\"Running CUDA Mode:\", device, torch.cuda.get_device_name(0))\n",
    "elif torch.backends.mps.is_available():\n",
    "    device = torch.device(\"mps\")\n",
    "    print(\"Running MPS Mode:\", device)\n",
    "else:\n",
    "    device = torch.device(\"cpu\")\n",
    "    print(\"Running CPU Mode:\", device)\n",
    "\n",
    "import time\n",
    "import math\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.ticker as ticker\n",
    "import numpy as np\n",
    "from tqdm import tqdm\n",
    "\n",
    "plt.switch_backend('agg')\n",
    "def asMinutes(s):\n",
    "    m = math.floor(s / 60)\n",
    "    s -= m * 60\n",
    "    return '%dm %ds' % (m, s)\n",
    "def timeSince(since, percent):\n",
    "    now = time.time()\n",
    "    s = now - since\n",
    "    es = s / (percent)\n",
    "    rs = es - s\n",
    "    return '%s (- %s)' % (asMinutes(s), asMinutes(rs))\n",
    "    \n",
    "def saveModel(save_path, model_state, optimiser_state, loss):\n",
    "    torch.save({\n",
    "            'model_state_dict': model_state,\n",
    "            'optimizer_state_dict':optimiser_state,\n",
    "            'loss': loss,  \n",
    "    }, save_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Class Definitions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Vocab and Dataloader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {},
   "outputs": [],
   "source": [
    "VOCAB = None\n",
    "MAX_SEQ_LEN = 128"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {},
   "outputs": [],
   "source": [
    "START_TOKEN = \"<START>\"\n",
    "END_TOKEN = \"<END>\"\n",
    "UNK_TOKEN = \"<UNK>\"\n",
    "PAD_TOKEN = \"<PAD>\"\n",
    "\n",
    "class Vocabulary:\n",
    "    def __init__(self, freq_dict, wd_to_id, id_to_wd):\n",
    "        self.freq_dict = freq_dict\n",
    "        self.wd_to_id = wd_to_id\n",
    "        self.id_to_wd = id_to_wd\n",
    "        self.N = len(freq_dict)\n",
    "    \n",
    "    def get_id(self, word):\n",
    "        if word in self.wd_to_id:\n",
    "            return self.wd_to_id[word]\n",
    "        else:\n",
    "            return self.wd_to_id[UNK_TOKEN]\n",
    "\n",
    "class LatexFormulaDataset(Dataset):\n",
    "    \"\"\"Latex Formula Dataset: Image and Text\"\"\"\n",
    "    \n",
    "    def __init__(self, csv_file, root_dir, transform = None, max_examples = None):\n",
    "        \"\"\"\n",
    "        Arguments:\n",
    "            csv_file (string): Path to the csv file with image name and text\n",
    "            root_dir (string): Directory with all the images.\n",
    "            transform (callable, optional): Optional transform to be applied\n",
    "                on a sample.\n",
    "        \"\"\"\n",
    "        info(\"Loading Dataset...\")\n",
    "        self.df = pd.read_csv(csv_file)\n",
    "        \n",
    "        #info(\"Loaded Dataset\", self.df.info)\n",
    "        \n",
    "        #Slice the dataset if max_examples is not None\n",
    "        if max_examples is not None:\n",
    "            self.df = self.df.iloc[:max_examples, :]\n",
    "\n",
    "        self.root_dir = root_dir\n",
    "        self.transform = transform\n",
    "\n",
    "        self.df['formula'] = self.df['formula'].apply(lambda x: x.split())\n",
    "        self.df['formula'] = self.df['formula'].apply(lambda x: [START_TOKEN] + x + [END_TOKEN])\n",
    "\n",
    "        self.maxlen = 0\n",
    "        for formula in self.df['formula']:\n",
    "            if len(formula) > self.maxlen:\n",
    "                self.maxlen = len(formula)\n",
    "     \n",
    "        \n",
    "        self.df['formula'] = self.df['formula'].apply(lambda x: x +[PAD_TOKEN]*(self.maxlen - len(x)))\n",
    "        self.vocab= self.construct_vocab() \n",
    "        # self.df['formula'] = self.df['formula'].apply(convert_to_ids)\n",
    "        info(\"Loaded.\")\n",
    "        \n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.df)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        \"\"\"\n",
    "        Returns sample of type image, textformula\n",
    "        \"\"\"\n",
    "        if torch.is_tensor(idx):\n",
    "            idx = idx.tolist()\n",
    "\n",
    "        img_name = os.path.join(self.root_dir,\n",
    "                                self.df.iloc[idx, 0])\n",
    "        image = io.imread(img_name)\n",
    "        formula = self.df.iloc[idx, 1]\n",
    "\n",
    "        # formula = np.array([formula], dtype=str).reshape(-1, 1)\n",
    "        # formula = [self.vocab.get_id(wd[0]) for wd in formula]\n",
    "        \n",
    "        def convert_to_ids(formula):\n",
    "            form2 = [VOCAB.get_id(wd) for wd in formula]\n",
    "            return torch.tensor(form2, dtype=torch.int64)[0:MAX_SEQ_LEN]\n",
    "        \n",
    "        sample = {'image': image, 'formula': convert_to_ids(formula)}\n",
    "\n",
    "        if self.transform:\n",
    "            sample['image'] = self.transform(sample['image'])\n",
    "            \n",
    "        return sample \n",
    "    \n",
    "    def construct_vocab(self):\n",
    "        \"\"\"\n",
    "        Constructs vocabulary from the dataset formulas\n",
    "        \"\"\"\n",
    "        #Split on spaces to tokenize\n",
    "        freq_dict = {}\n",
    "        for formula in self.df['formula']:\n",
    "            for wd in formula:\n",
    "                if wd not in freq_dict:\n",
    "                    freq_dict[wd] = 1\n",
    "                else:\n",
    "                    freq_dict[wd] += 1\n",
    "        freq_dict[UNK_TOKEN] = 1\n",
    "        N = len(freq_dict)\n",
    "        wd_to_id = {}\n",
    "        for i, wd in enumerate(freq_dict):\n",
    "            wd_to_id[wd] = i\n",
    "        id_to_wd = {v: k for k, v in wd_to_id.items()}\n",
    "    \n",
    "        #pad the formulas with \n",
    "        return Vocabulary(freq_dict, wd_to_id, id_to_wd)      \n",
    "\n",
    "def get_dataloader(csv_path, image_root, batch_size, transform = None, max_examples = None):\n",
    "    \"\"\"\n",
    "    Returns dataloader,dataset for the dataset\n",
    "    \"\"\"\n",
    "    dataset = LatexFormulaDataset(csv_path, image_root, max_examples=max_examples,transform=transform) #checked\n",
    "    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)\n",
    "    return dataloader, dataset\n",
    " "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Encoder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {},
   "outputs": [],
   "source": [
    "#import vgg16 from torch\n",
    "from torchvision.models import vgg16"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Decoder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DecoderLSTM(nn.Module):\n",
    "    \"\"\"\n",
    "    Inputs:\n",
    "    (here M is whatever the batch size is passed)\n",
    "\n",
    "    context_size : size of the context vector [shape: (1,M,context_size)]\n",
    "    n_layers: number of layers [for our purposes, defaults to 1]\n",
    "    hidden_size : size of the hidden state vectors [shape: (n_layers,M,hidden_size)]\n",
    "    embed_size : size of the embedding vectors [shape: (1,M,embed_size)]\n",
    "    vocab_size : size of the vocabulary\n",
    "    max_length : maximum length of the formula\n",
    "    \"\"\"\n",
    "    def __init__(self, context_size, vocab, max_seq_len, n_layers = 1, hidden_size = 512, embed_size = 512):\n",
    "        super().__init__()\n",
    "        self.context_size = context_size\n",
    "        self.vocab = vocab\n",
    "        self.vocab_size = vocab.N\n",
    "        self.n_layers = n_layers\n",
    "        self.hidden_size = hidden_size\n",
    "        self.embed_size = embed_size\n",
    "        self.input_size = context_size + embed_size\n",
    "        self.max_seq_len = max_seq_len\n",
    "\n",
    "        self.embed = nn.Embedding(self.vocab_size, embed_size)\n",
    "        self.lstm = nn.LSTMCell(self.input_size, self.hidden_size)\n",
    "        self.linear = nn.Linear(hidden_size, self.vocab_size)\n",
    "    \n",
    "    def forward(self, context, target_tensor = None):\n",
    "        \"\"\"\n",
    "        M: batch_size\n",
    "        context is the context vector from the encoder [shape: (M,context_size)]\n",
    "        target_tensor is the formula in tensor form [shape: (M,max_length)] (in the second dimension, it is sequence of indices of formula tokens)\n",
    "            if target_tensor is not None, then we are in Teacher Forcing mode\n",
    "            else normal jo bhi (last prediction ka embed is concatenated)\n",
    "        \"\"\"\n",
    "        context.to(device)\n",
    "        batch_size = context.shape[0]\n",
    "\n",
    "        #initialize hidden state and cell state\n",
    "        hidden = context\n",
    "        cell = torch.zeros(batch_size, self.hidden_size).to(device)\n",
    "\n",
    "        #initialize the input with embedding of the start token. Expand for batch size.\n",
    "        init_embed = self.embed(torch.tensor([self.vocab.wd_to_id[START_TOKEN]]).to(device).expand(batch_size, -1)).squeeze()\n",
    "        \n",
    "        #initialize the output_history and init_output\n",
    "        outputs = []\n",
    "        output = torch.zeros((batch_size, self.vocab_size)).to(device)\n",
    "        \n",
    "        \n",
    "        for i in range(self.max_seq_len):\n",
    "            #teacher forcing: 50% times\n",
    "            r = torch.rand(1)\n",
    "            if r>0.5 and target_tensor is not None:\n",
    "                if i==0 :\n",
    "                    embedding = init_embed\n",
    "                else: \n",
    "                    embedding = self.embed(target_tensor[:, i-1]).reshape((batch_size, self.embed_size)).to(device)            \n",
    "            else:\n",
    "                if i==0 :\n",
    "                    embedding = init_embed\n",
    "\n",
    "                else:\n",
    "                    #create embedding from previous input\n",
    "                    embedding = self.embed(torch.argmax(output, dim = 1))\n",
    "\n",
    "            lstm_input = torch.cat([context, embedding], dim = 1).to(device)\n",
    "    \n",
    "            hidden, cell = self.lstm(lstm_input, (hidden, cell))\n",
    "            output = self.linear(hidden)\n",
    "            outputs.append(output)\n",
    "            \n",
    "        output_tensor = torch.stack(outputs).permute(1,0,2) #LBV - > BLV\n",
    "\n",
    "        return output_tensor, hidden, cell\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Final"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {},
   "outputs": [],
   "source": [
    "class HandwritingToLatexModel(nn.Module):\n",
    "    def __init__(self, context_size, vocab, max_seq_len, n_layers, hidden_size, embed_size):\n",
    "        super().__init__()\n",
    "        self.encoder = vgg16(weights=\"IMAGENET1K_V1\")\n",
    "        #Freeze weights of the encoder and unfreeze last layer\n",
    "        self.encoder.classifier[-1] = nn.Linear(4096, 512)\n",
    "        for param in self.encoder.parameters():\n",
    "            param.requires_grad = False\n",
    "        for param in self.encoder.classifier[-1].parameters():\n",
    "            param.requires_grad = True\n",
    "        self.decoder = DecoderLSTM(context_size, vocab, max_seq_len, n_layers, hidden_size, embed_size)\n",
    "    \n",
    "    def forward(self, image, target_tensor = None):\n",
    "        context = self.encoder(image)\n",
    "        outputs, _, _ = self.decoder(context, target_tensor)\n",
    "        return outputs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {},
   "outputs": [],
   "source": [
    "DATA_BASE_PATH = \"data/\"\n",
    "SAVE_BASE_PATH = \"checkpoints/\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training Code"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_epoch(dataloader,model, optimizer, criterion):\n",
    "    total_loss = 0\n",
    "    idx = 0\n",
    "    pb = tqdm(dataloader, desc=\"Batch\")\n",
    "    for data in pb:\n",
    "        idx+=1\n",
    "        input_tensor, target_tensor = data['image'].to(device), data['formula'].to(device)\n",
    "        outputs = model(input_tensor, target_tensor)\n",
    "        train_dataset = dataloader.dataset \n",
    "        if(train_dataset and idx%100==0):\n",
    "            generated_formula = [VOCAB.id_to_wd[token.item()] for token in torch.argmax(outputs, dim=2)[0]]\n",
    "            required_formula = [VOCAB.id_to_wd[token.item()] for token in target_tensor[0]]\n",
    "            print(f\"Generated: {' '.join(generated_formula)}\")\n",
    "            print(f\"Actual: {' '.join(required_formula)}\")\n",
    "\n",
    "        output_logits = outputs.permute(0,2,1)\n",
    "        \n",
    "        loss = criterion(\n",
    "            output_logits,\n",
    "            target_tensor\n",
    "        )\n",
    "        \n",
    "        #backprop\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        optimizer.step() \n",
    "\n",
    "        total_loss += loss.item()\n",
    "        pb.set_description(f\"Loss: {loss.item()}\")\n",
    "\n",
    "    return total_loss / len(dataloader)\n",
    "\n",
    "def train(train_dataloader, model, n_epochs, optimizer = None, learning_rate=0.001, print_every=1, save_interval=2, save_prefix = 'model'):\n",
    "    start = time.time()\n",
    "    plot_losses = []\n",
    "    print_loss_total = 0  # Reset every print_every\n",
    "\n",
    "    if not optimizer: optimizer = optim.Adam(model.parameters(), lr=learning_rate)\n",
    "    criterion = nn.CrossEntropyLoss(ignore_index=VOCAB.wd_to_id[PAD_TOKEN]).to(device) #as stated in assignment\n",
    "\n",
    "    model.train()\n",
    "    for epoch in range(1, n_epochs + 1):\n",
    "        info(f\"Epoch {epoch}\")\n",
    "        loss = train_epoch(train_dataloader, model, optimizer, criterion)\n",
    "        print_loss_total += loss\n",
    "\n",
    "        if epoch % print_every == 0:\n",
    "            print_loss_avg = print_loss_total / print_every\n",
    "            print_loss_total = 0\n",
    "    \n",
    "        if epoch % save_interval == 0:\n",
    "            saveModel(f'{SAVE_BASE_PATH}{save_prefix}_epoch_{epoch}.pt', model.state_dict(), optimizer.state_dict(), loss)    \n",
    "                \n",
    "        info('%s (%d %d%%) %.4f' % (timeSince(start, epoch / n_epochs), epoch, epoch / n_epochs * 100, print_loss_avg))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load Datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading Dataset...\n",
      "Loaded.\n",
      "Loading Dataset...\n",
      "Loaded.\n"
     ]
    }
   ],
   "source": [
    "#part a\n",
    "batch_size = 32\n",
    "vocab_size = 1000\n",
    "CONTEXT_SIZE = 512\n",
    "HIDDEN_SIZE = 512\n",
    "EMBED_SIZE = 512\n",
    "MAX_EXAMPLES = 1000\n",
    "\n",
    "transform = transforms.Compose([\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Resize((224, 224)),\n",
    "])\n",
    "\n",
    "# transform = transforms.Compose([\n",
    "#     transforms.ToTensor(),\n",
    "#     transforms.Resize((224, 224)),\n",
    "#     transforms.Normalize((0.48235, 0.45882, 0.40784), (0.00392156862745098, 0.00392156862745098, 0.00392156862745098))\n",
    "# ])\n",
    "\n",
    "train_csv_path = f\"{DATA_BASE_PATH}/SyntheticData/train.csv\"\n",
    "image_root_path = f\"{DATA_BASE_PATH}/SyntheticData/images/\"\n",
    "train_dataloader, train_dataset = get_dataloader(train_csv_path, image_root_path, batch_size, transform, max_examples=None)\n",
    "\n",
    "hw_train_csv_path = f\"{DATA_BASE_PATH}/HandwrittenData/train_hw.csv\"\n",
    "hw_image_root_path = f\"{DATA_BASE_PATH}/HandwrittenData/images/train/\"\n",
    "hw_train_dataloader, hw_train_dataset = get_dataloader(hw_train_csv_path, hw_image_root_path, batch_size, transform, max_examples=None)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Create vocabulary connsisting of both vocabs of synthetic and handwritten datasets\n",
    "def combine_vocabs(v1, v2):\n",
    "    freq_dict = {}\n",
    "    for wd in v1.freq_dict:\n",
    "        freq_dict[wd] = v1.freq_dict[wd]\n",
    "    for wd in v2.freq_dict:\n",
    "        if wd not in freq_dict:\n",
    "            freq_dict[wd] = v2.freq_dict[wd]\n",
    "        else:\n",
    "            freq_dict[wd] += v2.freq_dict[wd]\n",
    "    freq_dict[UNK_TOKEN] = 1\n",
    "    N = len(freq_dict)\n",
    "    wd_to_id = {}\n",
    "    for i, wd in enumerate(freq_dict):\n",
    "        wd_to_id[wd] = i\n",
    "    id_to_wd = {v: k for k, v in wd_to_id.items()}\n",
    "    return Vocabulary(freq_dict, wd_to_id, id_to_wd)\n",
    "\n",
    "VOCAB = combine_vocabs(train_dataset.vocab, hw_train_dataset.vocab)\n",
    "\n",
    "model = HandwritingToLatexModel(CONTEXT_SIZE, VOCAB, n_layers=1, hidden_size= HIDDEN_SIZE, embed_size=EMBED_SIZE, max_seq_len=MAX_SEQ_LEN).to(device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loss: 3.409292221069336:   4%|▍         | 99/2344 [01:22<29:55,  1.25it/s] "
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Generated: <START> $ { { } } } } { { } { { } } } } } { { } } } } } } { { } } { { } } } } { { } } { { } } { $ <END> { { { { { { } } { { { { { { { { { } { } { } } { { { { { } { } } { { { { } } { { } } { { { { } } { } } { { { { } } { { } } { { { { { { { } } { { { } } { { { { } }\n",
      "Actual: <START> $ T _ { { \\bf \\L } } \\equiv 2 \\mathrm { i } \\mathrm { I m } { \\cal N } _ { { \\bf \\L } { \\bf \\S } } L ^ { { \\bf \\S } } , $ <END> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD>\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loss: 3.058997631072998:   8%|▊         | 199/2344 [02:43<28:25,  1.26it/s] "
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Generated: <START> $ { _ { { 2 } } { { 2 } } { { { 2 } } { 2 1 } { 2 } } { { } } { 2 } { 2 } } { { } } { { x ) { 2 } } \\frac { 2 } } } } { 2 } } { { } } { { 2 } { 2 } } { 2 { { { { { 2 { 2 { 2 } } { { { 2 { 2 { { { 2 { 2 { 2 } } { { { { { 2 { 2 } } { { { { 2 { { 2 } { { { { 2\n",
      "Actual: <START> $ \\partial \\phi _ { 0 } \\gamma _ { 2 1 } \\phi _ { 0 } ^ { - 1 } = \\frac { \\mathrm { e } ^ { \\gamma _ { 5 } \\beta } } { \\rho } ( j _ { 0 } + \\gamma _ { 5 } g _ { 0 } ) . $ <END> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD> <PAD>\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loss: 2.965256929397583:  11%|█         | 255/2344 [03:30<28:42,  1.21it/s] \n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "\u001b[1;32m/Users/samarth/Library/CloudStorage/OneDrive-IITDelhi/Sem5/COL774/assgns/assgn4/Handwriting-to-Latex-ML/main-vgg.ipynb Cell 26\u001b[0m line \u001b[0;36m1\n\u001b[0;32m----> <a href='vscode-notebook-cell:/Users/samarth/Library/CloudStorage/OneDrive-IITDelhi/Sem5/COL774/assgns/assgn4/Handwriting-to-Latex-ML/main-vgg.ipynb#X42sZmlsZQ%3D%3D?line=0'>1</a>\u001b[0m train(train_dataloader, model, \u001b[39m25\u001b[39;49m, save_interval\u001b[39m=\u001b[39;49m\u001b[39m2\u001b[39;49m, save_prefix \u001b[39m=\u001b[39;49m \u001b[39m'\u001b[39;49m\u001b[39mmodel4.0\u001b[39;49m\u001b[39m'\u001b[39;49m)\n",
      "\u001b[1;32m/Users/samarth/Library/CloudStorage/OneDrive-IITDelhi/Sem5/COL774/assgns/assgn4/Handwriting-to-Latex-ML/main-vgg.ipynb Cell 26\u001b[0m line \u001b[0;36m4\n\u001b[1;32m     <a href='vscode-notebook-cell:/Users/samarth/Library/CloudStorage/OneDrive-IITDelhi/Sem5/COL774/assgns/assgn4/Handwriting-to-Latex-ML/main-vgg.ipynb#X42sZmlsZQ%3D%3D?line=41'>42</a>\u001b[0m \u001b[39mfor\u001b[39;00m epoch \u001b[39min\u001b[39;00m \u001b[39mrange\u001b[39m(\u001b[39m1\u001b[39m, n_epochs \u001b[39m+\u001b[39m \u001b[39m1\u001b[39m):\n\u001b[1;32m     <a href='vscode-notebook-cell:/Users/samarth/Library/CloudStorage/OneDrive-IITDelhi/Sem5/COL774/assgns/assgn4/Handwriting-to-Latex-ML/main-vgg.ipynb#X42sZmlsZQ%3D%3D?line=42'>43</a>\u001b[0m     info(\u001b[39mf\u001b[39m\u001b[39m\"\u001b[39m\u001b[39mEpoch \u001b[39m\u001b[39m{\u001b[39;00mepoch\u001b[39m}\u001b[39;00m\u001b[39m\"\u001b[39m)\n\u001b[0;32m---> <a href='vscode-notebook-cell:/Users/samarth/Library/CloudStorage/OneDrive-IITDelhi/Sem5/COL774/assgns/assgn4/Handwriting-to-Latex-ML/main-vgg.ipynb#X42sZmlsZQ%3D%3D?line=43'>44</a>\u001b[0m     loss \u001b[39m=\u001b[39m train_epoch(train_dataloader, model, optimizer, criterion)\n\u001b[1;32m     <a href='vscode-notebook-cell:/Users/samarth/Library/CloudStorage/OneDrive-IITDelhi/Sem5/COL774/assgns/assgn4/Handwriting-to-Latex-ML/main-vgg.ipynb#X42sZmlsZQ%3D%3D?line=44'>45</a>\u001b[0m     print_loss_total \u001b[39m+\u001b[39m\u001b[39m=\u001b[39m loss\n\u001b[1;32m     <a href='vscode-notebook-cell:/Users/samarth/Library/CloudStorage/OneDrive-IITDelhi/Sem5/COL774/assgns/assgn4/Handwriting-to-Latex-ML/main-vgg.ipynb#X42sZmlsZQ%3D%3D?line=46'>47</a>\u001b[0m     \u001b[39mif\u001b[39;00m epoch \u001b[39m%\u001b[39m print_every \u001b[39m==\u001b[39m \u001b[39m0\u001b[39m:\n",
      "\u001b[1;32m/Users/samarth/Library/CloudStorage/OneDrive-IITDelhi/Sem5/COL774/assgns/assgn4/Handwriting-to-Latex-ML/main-vgg.ipynb Cell 26\u001b[0m line \u001b[0;36m8\n\u001b[1;32m      <a href='vscode-notebook-cell:/Users/samarth/Library/CloudStorage/OneDrive-IITDelhi/Sem5/COL774/assgns/assgn4/Handwriting-to-Latex-ML/main-vgg.ipynb#X42sZmlsZQ%3D%3D?line=5'>6</a>\u001b[0m idx\u001b[39m+\u001b[39m\u001b[39m=\u001b[39m\u001b[39m1\u001b[39m\n\u001b[1;32m      <a href='vscode-notebook-cell:/Users/samarth/Library/CloudStorage/OneDrive-IITDelhi/Sem5/COL774/assgns/assgn4/Handwriting-to-Latex-ML/main-vgg.ipynb#X42sZmlsZQ%3D%3D?line=6'>7</a>\u001b[0m input_tensor, target_tensor \u001b[39m=\u001b[39m data[\u001b[39m'\u001b[39m\u001b[39mimage\u001b[39m\u001b[39m'\u001b[39m]\u001b[39m.\u001b[39mto(device), data[\u001b[39m'\u001b[39m\u001b[39mformula\u001b[39m\u001b[39m'\u001b[39m]\u001b[39m.\u001b[39mto(device)\n\u001b[0;32m----> <a href='vscode-notebook-cell:/Users/samarth/Library/CloudStorage/OneDrive-IITDelhi/Sem5/COL774/assgns/assgn4/Handwriting-to-Latex-ML/main-vgg.ipynb#X42sZmlsZQ%3D%3D?line=7'>8</a>\u001b[0m outputs \u001b[39m=\u001b[39m model(input_tensor, target_tensor)\n\u001b[1;32m      <a href='vscode-notebook-cell:/Users/samarth/Library/CloudStorage/OneDrive-IITDelhi/Sem5/COL774/assgns/assgn4/Handwriting-to-Latex-ML/main-vgg.ipynb#X42sZmlsZQ%3D%3D?line=8'>9</a>\u001b[0m train_dataset \u001b[39m=\u001b[39m dataloader\u001b[39m.\u001b[39mdataset \n\u001b[1;32m     <a href='vscode-notebook-cell:/Users/samarth/Library/CloudStorage/OneDrive-IITDelhi/Sem5/COL774/assgns/assgn4/Handwriting-to-Latex-ML/main-vgg.ipynb#X42sZmlsZQ%3D%3D?line=9'>10</a>\u001b[0m \u001b[39mif\u001b[39;00m(train_dataset \u001b[39mand\u001b[39;00m idx\u001b[39m%\u001b[39m\u001b[39m100\u001b[39m\u001b[39m==\u001b[39m\u001b[39m0\u001b[39m):\n",
      "File \u001b[0;32m~/.pyenv/versions/3.10.4/lib/python3.10/site-packages/torch/nn/modules/module.py:1518\u001b[0m, in \u001b[0;36mModule._wrapped_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1516\u001b[0m     \u001b[39mreturn\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_compiled_call_impl(\u001b[39m*\u001b[39margs, \u001b[39m*\u001b[39m\u001b[39m*\u001b[39mkwargs)  \u001b[39m# type: ignore[misc]\u001b[39;00m\n\u001b[1;32m   1517\u001b[0m \u001b[39melse\u001b[39;00m:\n\u001b[0;32m-> 1518\u001b[0m     \u001b[39mreturn\u001b[39;00m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49m_call_impl(\u001b[39m*\u001b[39;49margs, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n",
      "File \u001b[0;32m~/.pyenv/versions/3.10.4/lib/python3.10/site-packages/torch/nn/modules/module.py:1527\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1522\u001b[0m \u001b[39m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m   1523\u001b[0m \u001b[39m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m   1524\u001b[0m \u001b[39mif\u001b[39;00m \u001b[39mnot\u001b[39;00m (\u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_backward_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_backward_pre_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_forward_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_forward_pre_hooks\n\u001b[1;32m   1525\u001b[0m         \u001b[39mor\u001b[39;00m _global_backward_pre_hooks \u001b[39mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m   1526\u001b[0m         \u001b[39mor\u001b[39;00m _global_forward_hooks \u001b[39mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1527\u001b[0m     \u001b[39mreturn\u001b[39;00m forward_call(\u001b[39m*\u001b[39;49margs, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n\u001b[1;32m   1529\u001b[0m \u001b[39mtry\u001b[39;00m:\n\u001b[1;32m   1530\u001b[0m     result \u001b[39m=\u001b[39m \u001b[39mNone\u001b[39;00m\n",
      "\u001b[1;32m/Users/samarth/Library/CloudStorage/OneDrive-IITDelhi/Sem5/COL774/assgns/assgn4/Handwriting-to-Latex-ML/main-vgg.ipynb Cell 26\u001b[0m line \u001b[0;36m1\n\u001b[1;32m     <a href='vscode-notebook-cell:/Users/samarth/Library/CloudStorage/OneDrive-IITDelhi/Sem5/COL774/assgns/assgn4/Handwriting-to-Latex-ML/main-vgg.ipynb#X42sZmlsZQ%3D%3D?line=12'>13</a>\u001b[0m \u001b[39mdef\u001b[39;00m \u001b[39mforward\u001b[39m(\u001b[39mself\u001b[39m, image, target_tensor \u001b[39m=\u001b[39m \u001b[39mNone\u001b[39;00m):\n\u001b[1;32m     <a href='vscode-notebook-cell:/Users/samarth/Library/CloudStorage/OneDrive-IITDelhi/Sem5/COL774/assgns/assgn4/Handwriting-to-Latex-ML/main-vgg.ipynb#X42sZmlsZQ%3D%3D?line=13'>14</a>\u001b[0m     context \u001b[39m=\u001b[39m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mencoder(image)\n\u001b[0;32m---> <a href='vscode-notebook-cell:/Users/samarth/Library/CloudStorage/OneDrive-IITDelhi/Sem5/COL774/assgns/assgn4/Handwriting-to-Latex-ML/main-vgg.ipynb#X42sZmlsZQ%3D%3D?line=14'>15</a>\u001b[0m     outputs, _, _ \u001b[39m=\u001b[39m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mdecoder(context, target_tensor)\n\u001b[1;32m     <a href='vscode-notebook-cell:/Users/samarth/Library/CloudStorage/OneDrive-IITDelhi/Sem5/COL774/assgns/assgn4/Handwriting-to-Latex-ML/main-vgg.ipynb#X42sZmlsZQ%3D%3D?line=15'>16</a>\u001b[0m     \u001b[39mreturn\u001b[39;00m outputs\n",
      "File \u001b[0;32m~/.pyenv/versions/3.10.4/lib/python3.10/site-packages/torch/nn/modules/module.py:1518\u001b[0m, in \u001b[0;36mModule._wrapped_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1516\u001b[0m     \u001b[39mreturn\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_compiled_call_impl(\u001b[39m*\u001b[39margs, \u001b[39m*\u001b[39m\u001b[39m*\u001b[39mkwargs)  \u001b[39m# type: ignore[misc]\u001b[39;00m\n\u001b[1;32m   1517\u001b[0m \u001b[39melse\u001b[39;00m:\n\u001b[0;32m-> 1518\u001b[0m     \u001b[39mreturn\u001b[39;00m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49m_call_impl(\u001b[39m*\u001b[39;49margs, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n",
      "File \u001b[0;32m~/.pyenv/versions/3.10.4/lib/python3.10/site-packages/torch/nn/modules/module.py:1527\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1522\u001b[0m \u001b[39m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m   1523\u001b[0m \u001b[39m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m   1524\u001b[0m \u001b[39mif\u001b[39;00m \u001b[39mnot\u001b[39;00m (\u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_backward_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_backward_pre_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_forward_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_forward_pre_hooks\n\u001b[1;32m   1525\u001b[0m         \u001b[39mor\u001b[39;00m _global_backward_pre_hooks \u001b[39mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m   1526\u001b[0m         \u001b[39mor\u001b[39;00m _global_forward_hooks \u001b[39mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1527\u001b[0m     \u001b[39mreturn\u001b[39;00m forward_call(\u001b[39m*\u001b[39;49margs, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n\u001b[1;32m   1529\u001b[0m \u001b[39mtry\u001b[39;00m:\n\u001b[1;32m   1530\u001b[0m     result \u001b[39m=\u001b[39m \u001b[39mNone\u001b[39;00m\n",
      "\u001b[1;32m/Users/samarth/Library/CloudStorage/OneDrive-IITDelhi/Sem5/COL774/assgns/assgn4/Handwriting-to-Latex-ML/main-vgg.ipynb Cell 26\u001b[0m line \u001b[0;36m6\n\u001b[1;32m     <a href='vscode-notebook-cell:/Users/samarth/Library/CloudStorage/OneDrive-IITDelhi/Sem5/COL774/assgns/assgn4/Handwriting-to-Latex-ML/main-vgg.ipynb#X42sZmlsZQ%3D%3D?line=64'>65</a>\u001b[0m         embedding \u001b[39m=\u001b[39m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39membed(torch\u001b[39m.\u001b[39margmax(output, dim \u001b[39m=\u001b[39m \u001b[39m1\u001b[39m))\n\u001b[1;32m     <a href='vscode-notebook-cell:/Users/samarth/Library/CloudStorage/OneDrive-IITDelhi/Sem5/COL774/assgns/assgn4/Handwriting-to-Latex-ML/main-vgg.ipynb#X42sZmlsZQ%3D%3D?line=66'>67</a>\u001b[0m lstm_input \u001b[39m=\u001b[39m torch\u001b[39m.\u001b[39mcat([context, embedding], dim \u001b[39m=\u001b[39m \u001b[39m1\u001b[39m)\u001b[39m.\u001b[39mto(device)\n\u001b[0;32m---> <a href='vscode-notebook-cell:/Users/samarth/Library/CloudStorage/OneDrive-IITDelhi/Sem5/COL774/assgns/assgn4/Handwriting-to-Latex-ML/main-vgg.ipynb#X42sZmlsZQ%3D%3D?line=68'>69</a>\u001b[0m hidden, cell \u001b[39m=\u001b[39m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mlstm(lstm_input, (hidden, cell))\n\u001b[1;32m     <a href='vscode-notebook-cell:/Users/samarth/Library/CloudStorage/OneDrive-IITDelhi/Sem5/COL774/assgns/assgn4/Handwriting-to-Latex-ML/main-vgg.ipynb#X42sZmlsZQ%3D%3D?line=69'>70</a>\u001b[0m output \u001b[39m=\u001b[39m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mlinear(hidden)\n\u001b[1;32m     <a href='vscode-notebook-cell:/Users/samarth/Library/CloudStorage/OneDrive-IITDelhi/Sem5/COL774/assgns/assgn4/Handwriting-to-Latex-ML/main-vgg.ipynb#X42sZmlsZQ%3D%3D?line=70'>71</a>\u001b[0m outputs\u001b[39m.\u001b[39mappend(output)\n",
      "File \u001b[0;32m~/.pyenv/versions/3.10.4/lib/python3.10/site-packages/torch/nn/modules/module.py:1518\u001b[0m, in \u001b[0;36mModule._wrapped_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1516\u001b[0m     \u001b[39mreturn\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_compiled_call_impl(\u001b[39m*\u001b[39margs, \u001b[39m*\u001b[39m\u001b[39m*\u001b[39mkwargs)  \u001b[39m# type: ignore[misc]\u001b[39;00m\n\u001b[1;32m   1517\u001b[0m \u001b[39melse\u001b[39;00m:\n\u001b[0;32m-> 1518\u001b[0m     \u001b[39mreturn\u001b[39;00m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49m_call_impl(\u001b[39m*\u001b[39;49margs, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n",
      "File \u001b[0;32m~/.pyenv/versions/3.10.4/lib/python3.10/site-packages/torch/nn/modules/module.py:1527\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1522\u001b[0m \u001b[39m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m   1523\u001b[0m \u001b[39m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m   1524\u001b[0m \u001b[39mif\u001b[39;00m \u001b[39mnot\u001b[39;00m (\u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_backward_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_backward_pre_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_forward_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_forward_pre_hooks\n\u001b[1;32m   1525\u001b[0m         \u001b[39mor\u001b[39;00m _global_backward_pre_hooks \u001b[39mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m   1526\u001b[0m         \u001b[39mor\u001b[39;00m _global_forward_hooks \u001b[39mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1527\u001b[0m     \u001b[39mreturn\u001b[39;00m forward_call(\u001b[39m*\u001b[39;49margs, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n\u001b[1;32m   1529\u001b[0m \u001b[39mtry\u001b[39;00m:\n\u001b[1;32m   1530\u001b[0m     result \u001b[39m=\u001b[39m \u001b[39mNone\u001b[39;00m\n",
      "File \u001b[0;32m~/.pyenv/versions/3.10.4/lib/python3.10/site-packages/torch/nn/modules/rnn.py:1347\u001b[0m, in \u001b[0;36mLSTMCell.forward\u001b[0;34m(self, input, hx)\u001b[0m\n\u001b[1;32m   1344\u001b[0m \u001b[39melse\u001b[39;00m:\n\u001b[1;32m   1345\u001b[0m     hx \u001b[39m=\u001b[39m (hx[\u001b[39m0\u001b[39m]\u001b[39m.\u001b[39munsqueeze(\u001b[39m0\u001b[39m), hx[\u001b[39m1\u001b[39m]\u001b[39m.\u001b[39munsqueeze(\u001b[39m0\u001b[39m)) \u001b[39mif\u001b[39;00m \u001b[39mnot\u001b[39;00m is_batched \u001b[39melse\u001b[39;00m hx\n\u001b[0;32m-> 1347\u001b[0m ret \u001b[39m=\u001b[39m _VF\u001b[39m.\u001b[39;49mlstm_cell(\n\u001b[1;32m   1348\u001b[0m     \u001b[39minput\u001b[39;49m, hx,\n\u001b[1;32m   1349\u001b[0m     \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mweight_ih, \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mweight_hh,\n\u001b[1;32m   1350\u001b[0m     \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mbias_ih, \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mbias_hh,\n\u001b[1;32m   1351\u001b[0m )\n\u001b[1;32m   1353\u001b[0m \u001b[39mif\u001b[39;00m \u001b[39mnot\u001b[39;00m is_batched:\n\u001b[1;32m   1354\u001b[0m     ret \u001b[39m=\u001b[39m (ret[\u001b[39m0\u001b[39m]\u001b[39m.\u001b[39msqueeze(\u001b[39m0\u001b[39m), ret[\u001b[39m1\u001b[39m]\u001b[39m.\u001b[39msqueeze(\u001b[39m0\u001b[39m))\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "train(train_dataloader, model, 25, save_interval=2, save_prefix = 'model4.0')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}