main.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 172,
   "id": "e8f2076e",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Global seed set to 42\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "42"
      ]
     },
     "execution_count": 172,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from importlib import reload\n",
    "import torch\n",
    "from torch import nn\n",
    "import pytorch_lightning as pl\n",
    "\n",
    "from pytorch_lightning import Trainer, seed_everything\n",
    "from pytorch_lightning.loggers import TensorBoardLogger\n",
    "from pytorch_lightning.callbacks import ModelCheckpoint, LearningRateMonitor\n",
    "from pytorch_lightning.callbacks.early_stopping import EarlyStopping\n",
    "\n",
    "seed_everything(42)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b82fc426",
   "metadata": {},
   "source": [
    "## Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 149,
   "id": "c956f167",
   "metadata": {},
   "outputs": [],
   "source": [
    "import src.tokenizer\n",
    "reload(src.tokenizer)\n",
    "from src.tokenizer import Tokenizer\n",
    "\n",
    "tokenizer = Tokenizer(alphabet=list(\" 'abcdefghijklmnopqrstuvwxyz\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 161,
   "id": "cfbe50f8",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torchaudio\n",
    "from torchaudio.transforms import MFCC\n",
    "from torch.utils.data import DataLoader, random_split\n",
    "\n",
    "mfcc = MFCC(sample_rate=16000, n_mfcc=40, melkwargs={'n_mels':64})\n",
    "def fn(batch):\n",
    "    X, _, Y, _, _, _  = list(zip(*batch))\n",
    "    X = [mfcc(x[0]).T for x in X]\n",
    "    Y = [torch.tensor(tokenizer.encode(y.lower())) for y in Y]\n",
    "    x_len = torch.tensor([x.shape[0] for x in X])\n",
    "    Mx = max(x_len)\n",
    "    y_len = torch.tensor([len(y) for y in Y])\n",
    "    My = max(y_len)\n",
    "    return {\n",
    "        'x' : nn.utils.rnn.pad_sequence(X, batch_first=True).transpose(1,2),\n",
    "        'x_len' : x_len,\n",
    "        'y' : nn.utils.rnn.pad_sequence(Y, batch_first=True),\n",
    "        'y_len': y_len\n",
    "    }\n",
    "\n",
    "dataset = torchaudio.datasets.LIBRISPEECH(root='data/', url='dev-clean', download=True)\n",
    "n = len(dataset)\n",
    "t = n // 10\n",
    "train_data, val_data = random_split(dataset, [n -t, t])\n",
    "tl = DataLoader(train_data, collate_fn=fn, batch_size=8)\n",
    "vl = DataLoader(val_data, collate_fn=fn, batch_size=8)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 162,
   "id": "8c3322f2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "x torch.Size([8, 40, 1694])\n",
      "x_len torch.Size([8])\n",
      "y torch.Size([8, 325])\n",
      "y_len torch.Size([8])\n"
     ]
    }
   ],
   "source": [
    "b = next(iter(tl))\n",
    "for k,v in b.items():\n",
    "    print(k, v.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 163,
   "id": "5111b160",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.image.AxesImage at 0x7f568dd41310>"
      ]
     },
     "execution_count": 163,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAABH4AAAA5CAYAAABAp+zeAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjUuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/YYfK9AAAACXBIWXMAAAsTAAALEwEAmpwYAAA26klEQVR4nO2dbdBl2VXX/2vvc869T3dPT89MT2bCZCQRAhIoEQoJVkpLgkJUimgVZQUpiIJGLYJgUSqJVWKJH1CRF6swViQxYKWIVIiawihGpLT8QHgJr0kIxISEiXmZSWb66X5e7j1n7+WHvdc+a+97nu7p6Z6nk6fXryp57j333HP22XufM73+97/WJmaGYRiGYRiGYRiGYRiGcfZwd7oBhmEYhmEYhmEYhmEYxrODCT+GYRiGYRiGYRiGYRhnFBN+DMMwDMMwDMMwDMMwzigm/BiGYRiGYRiGYRiGYZxRTPgxDMMwDMMwDMMwDMM4o3SnebLh3j1eP3zxNE+5CBFAYESm3c/y/+nFzgjAM137jEEgcPn7bPBsH1t4pud4Ntt3llnqe71Nti/172n1+Wfz2N7Kff3ZyPXG6rN5HA3DMAzDMAzDAK7+7iefYOYHlz47VeFn/fBFvPjf/JWd7UQMziJMVIGtOyEQ0fu379tjyTGI0l9HSfBhJkzsyjYHLvvLvgCq80jb5Jj6fQQVISkyoXOxarM+LjNV1+kpVudq26H3l3NP7Mq1SPuX+k8fS46j913a56Q+B4DArjq+pj3e9fqs9BPFnTHT/aGPJdd8vfNKf8jfG6HnRzsH235aOt/Ssdr9vIsI0VXH9hQR2FX9tTT3pV0nzVW93VMsfSZzsct9qefSSX2mz93O7XZcnw5EXK6x5aRrXupT2abbqcdm6R7Q52zntZ5jbZ+c1BY9BzuKJ4799bYv9UM7Lu123Ua979K9eqP52J5v6T6VOdTeD7qPT3re3gzP9HuGYRiGYRiGYSzzP176Ix8+6bNTFX6EnaCS50CgFUH0dyRIaYPJKnDLx+hcROQUsAR2cJw+D4zK6dOpALAE0owSYMq+J4kJS9sGFwDsClLSVi3apPPWQdhJQpDuOwk+l4QQuRaNHEsCu5P2WQpQfRaxWodUO45L5xbho/QVz9/tGoHHU4RHHRAuiRPtueR8+lwOtNj/7bWWY50QhBbRoREqno5AKWjRp+0bjR5jmSNgqubgznfytiQApfvCU4TLx9hGj87V/doKBH2er7qvZJzkdfsdDRFXgmq5bnap7bTcR7pPp5hFPSXs6PFx4Oo4MrYO83eKuNu2d2FspW0yB5fm2JIwKfufNOYnXefSawcu7dgRKZs26/t2SciS48lnS8JN78KOcKvnpLQHQHWfynWX8WAlBD4Dk5CJPoZhGIZhGIZxupy68HNSEFREhywCLTl3dLAnaMeDBFCBHUJwtVBBs7BT3BBuDurkHBLgTFEFSMQlMG1xxMUFAEoBmhZ1gDrQKe3NAevErrqGtj/kmvS1Xy8IP0lMKO6LhUDtJDGinEf1hRYh2mMsORIiCDH6neMvOV20A0a7V65HG/TK+XWAWm1T575eP+yIDjdow5IjZqkvtdCk+6JFu0r06zZoLsdQTiCft3UUsWWf5pC6HhE0ioOquLlop69aAWLJjcZM5btVH58g1pXPleDZubgrOGZXFDDfN+2YtM4d7XZqxVvZVtwrVDubvIuVqOxdLNe35IxqHUgnOfpKG2i+rnafkxxysq0SbBaEuDKneR63SjCU+3HB3aX70XG67vLd64h+JzncboTu45t1kRmGYRiGYRiGcfOcuvBzo5SinV+UoVJbeNf9UqV+0G7alaCFHGAOlFv3iw5EowqeBheqdLEQlWOA5msL7MA8B44x748ccHaoxaY2bUQEmvaatENBB2uB3U4wLu3TiGh0w75XfbGUvqIdUq0bphUMtBtnyXGjg019/V59R15H9lV7llJ1RHCi7JTRxwRQjiV9IQFu657Qc+56jhXdR0ss9ace+1Yo020DVDBOjF6JGIF3g2WdOqavWYSOCa4SPIDkROr8hBBd2X9J4Fi63hulGQH1XJF9pF/FSSRpcL0L1VjIfaDFCy2q6j6r3XNcizPYvRbt8tH7xeh2HD713KmdhlqQ1X2iHVNybC206b7zFOEJi88h/X5pvPW1L4nCxRWo5p12d2mBVYuuaBxqss/S8+AkThJD9T3diviGYRiGYRiGYTw7nLrwI4HUSXV03ELQD2DHqULE6HWaVnbldC5iig6Rk8ukBJ80uwZke+9TsCmBrxZHHGMWdCiJBE4FmgCAWLfTEcOjCURVwK6DVbnWk5wnsl9PsYhM6ZSq3gbP+7ZOhCV0asZJv7TrNsmxpZ9jE+Bez52jnUqayj2Bul5N69LRbdQBuW8EBaB2wmh3xNSkBcp3dbt0eoxOccJCqs9SOk05VtPvsm9gV0SL1L66H5eupThSuO5T/bmIkyJq6nQgaVdxoikXh3zmXazS8GQuiztoqa5Le99qcQBQoqhKldL9Ie3sXMQYkhspgkBK9JF5qkWciNlZRNkpJ+fV95aeh22KXNt2fe+IK0r6Iyw4fbRrR0RhGac2BU1/T/qmQ6zGwaMe43YeLrnZRKCRvhrFCbXgatKiVytOCXK8JYHppJTT1o10kniz9GxbSl/V+5oQZBiGYRiGYRi3n1MXfsboS0BU6kaoAK2817+Y85yaBehgfg42x+BBxNiE+ZLaX/4liJfPpuh2AiIReKbosMq/jstxJN1riq4EtxLsSwCq26sDZQn+Aey0UTsuSpsZ2R2ghBiaRRC5prTvXLtInAVtH7RpN+2v8ScFcBJsy+eTCjSZU6CshQJ9Lt0/OyloSgCUIFeLWx3F4gIZoy9Cl05lAlA+i1IjRoStJujX7h6ZP4EboUuLYvqz1u3As3Al/S+06WvyWWnnwlzX1ybHDNEhNgKQ7FvGRAkeTs19cZWJCBrYAXm+yhwMkp6oWCrErEWH1gkjY6cFlRCXUxcr8SzfQ3JOLeRpt1brfpP+0W1vawq1Y6/HRY+Tvq5N6Mo80U65VgjrXZjPR/W9JSKLFmz0eEnb5fmnKa4/nq9Puw1Traf6vtLPzirlTATjpm5T6QN53fSR1IKqjpe7T/pGC5maG6VrnSTyPB3nkGEYhmEYhmEYt84NhR8iehTATwJ4CCkUeD0z/ygR/WMAfwPA43nX1zLzO254QoqIVP/Df/7FfRYWtLjRCh1awJnbObsJqhoW+djE80pbU3QYVd0ZYP4Fv8sBnl6VSwdhU6xrjejgTjtvtJgDoPxKX4r2QgWm2k3UBFBSk0hqoLTBm+6bneOidgLoAHUpWFuqOaJX0qo+JwZo3kenvJW2R4fBhRSgNkHq7H6pg0g5dgRhCl3pxzpQrVOAdEpcS8x9DhGPxAUDVO3S9V7avpB+a9sI1HVoSnrYgvNJF+PW+y71rQgFfZ7Lui6NoJ1EaX9g5SdsQpfGA7HMVREBSAlrMYt2Or3qpFSeykGVBbYOKkWRlwUFPR5VoeEsfrTFvQEUAW9JNNTij74vypioPtXbTloBS85V2k2NW06JiMwE+Poa9Wek7oXqPNolpJ5f8n6KDhNUnTL1TGznCjCP5Rh9JVJJL8zzZFc80mOiReAIKvep3q99xi0Jcbo/r+dgLOeS56wu+H7Cs88wDMMwDMMwjNvD03H8TAC+h5nfTUT3APhVInpn/uyHmfkHb/akbbAMLLtZdFrBTpDBqIQWYA4uJR2EmRAbh4+kyAwUKudORxHez0GkBM1t2oYcQwI+2d77kNPDZpFCr9gjwaqkbUltk8iE3ofKfdQ6KMTRJPvqAK1Na/EuVs4VWcnnpLotSwG0fFeCtLZeUEm50qt9Eco16HHZhK64qLToJ0Fi70LlgtDBqu6L9rjSN1E5W/a6sWprmyqlj1mKBsdd4an0geozOW5Zmr1c+ywgts4sEaTkszZFTa9mVwXaue9FnNT1hsqxczeU6wgegwtpLmanlL6OIia5kGpQaccK5nnezsNItZOran8T0Os+u55YpvurTS1cEg+rPqVZHGamkt4JoCrcLvMzCUi752/nm3Yh6f1kzuprXRJb9bXr623nc+rYtG/vApzjHbFIjlXuXTWvinuIGM6HE8XbJVeljMnScvTitJLxaF2Dbcrkzbh1WsG47SO9n4k/hmEYhmEYhnH7uaHww8wfA/Cx/PoqEb0PwCPP5GQMwiZ0dV2N/HfJYdMGCRqd8lNcPTkw6n0owVnrmJmaNApJ52De/XUcQCmkW5YJd/Mv3+Ik0gLQGP2uO0ZqiYABF+GBkpLViygEByg3ilzjyk/lvaTilGtSxVt1HKbPLeKBHHcrYgLVwX+LFJ0FdlfpagN+GTtZxl63SQeQ7YpqMtbyV0S51gVQCQPqOJEJqxz8SsqcPtbSHErOsjpo1e3WaOdDp8bB+WbFKrmmPKY6ZU0fZ9R9X9pS3wfiXJp4txizON12HEaNg0LETHGHiDCGLHTKd7XYIWlbrQAo7dXLrWtRp4izCzpAK5SG6ND7On1yaYx3nGVYFlHKPamfA7TsWtPj1B5Txh90stNLRLHillOioj6eHhO9EmB7veVZoRw2S25ALRxrh1qZa41rUrMk9kamIua1orlug3YHyba2Ldc7T7t96X5p++Wk4xiGYRiGYRiGcWvcVI0fIno+gC8D8C4ALwHwaiL6VgC/guQKenLhO68C8CoAWD10zxz4qWBxwu6vyZJ6IFRBbg6Adb2gSijgtIqRDrw20+6ligMHQHLSxDnNTDssdN0NR1zXCiEuxXq3TfoYkF1IDgjBwbs4LxevBBH5XudiVf8HACaaXQja7RCZsOX5e5XwpIJk+Uz6dUkI0a9nEcJXgW2M9epX3tXuKBEXlgSXNtDTY5VeR7hsYUkOn12xSAS9Mj5U1yzSf/W55Dz6+nV/SyC/FIxK4K/TedoAvnVr6OtvaYWBEB1CvuZWRJBxC9HB+dqNsiQ46OsjSjWPigiZ50FQfeRQ19nRbi2grk91UjCu3UHihPKNgFsJSY2AJG3Wfd1RrO477VYRSuFqF6tj6DHRbril1cUqx08jhMgzRvpSrlEEIn2ft2OxJMTqaxmDL88B6UN97NKm/AzTQpu+r9o0wbZftVAk7ZE6aG1qpBaU9DVV17DgSGsFNJ0qdz0hSAuCGhN9DMMwDMMwDOPZ4WkLP0R0AcDPAPhuZt4notcB+H6kRIrvB/AvAXxb+z1mfj2A1wPAxS98iCWYbV0/2yxYANgNepqgXAqRdhRBbg5UNW0Kj05Lkn07inC+TjsQd44Ho3dz0CzBirgptAtIp05oF0UlXimRQQJmr8SSjtKKP70L1TmAObBdcuHIawlSASmWWwdW7S/0Gmm/9Gd7bABluw4QW5FDXFWyr4giVVDbODvEMSXCQVWMVp9Hxij3ia51Itcgx9ftkqBZOxh0m2ahwpW0M52GtiSSibuoPW67jLxGp2+JgCV/AVSiWZkylBxiydk0lW1a8CxtghawVO0apFo2zFQJFis/7aS7FQcNcbWik04P0qKMFhwdGJ2Pu/MmCyelH5vuaUW3SoCQtCrC4v3QOnrktaR4tueU/T3NQqMsc98WRZdzTGHuMxEttJha3FSY74cl14y0Xzv4RMCMvJtaJcdaStVs+6u81nOP1LYs8qy7cT5G9NUzV88lXeRZnn9L6WiL573B51qgWnI1GYZhGIZhGIZx+3lawg8R9Uiiz5uZ+W0AwMyfUJ//WwA/+3SOJcEssBsw6HQfeb/0C7JO2wLPzoLiWGlqfegAuw3E9HnLr9r5u2NMv9DPokpicCE5UeIsWuhrKJ9rNwPPzpGYvyMCQsiuA85un/bX8NZBol/rYN2x9GkdpMlqREt1ObTjQUSUdjlxnUIFZHdUI/7I5zpdRdqt0/v0cQCU1JzNOE/FpeXpKyGDCaIESb2otq+DCsZbR03lFFFt0fWXHO+uxKXR2yfOdWYa0aod71J3hmt3he7r9hwS+B9PfTVvZL4WMZGTcDXFbhYr1VzXK211atl3fS1akNXpZFIwnURUUk6nMczpa9ohVAQFFcu3gk1bE0mohAeax7MstZ5fj9FXgm45L3EtxrZ1wniuZVMtN9+0o3XMAbWgrGuISb0p3Z9jSCLfznExn18E5JPqOWmRuhVNdVva50Nsvic1esZYzzu9f3ut+pr1mLZ9tPRab9Mipcw/6YdK2FcuJMMwDMMwDMMwbh9PZ1UvAvAGAO9j5h9S25+b6/8AwF8C8NtP54RtwCFo+7+nWNW/AFBSJPQx9EpJEgQVUUjEj6WgtnEc6WO2yzmvuqkKdNu2OxUsSTu27Iu4IMH5ln1xrvQq1WPlp/nX9NymPtet0UKNFOzVAZS4Tsbg4XyoanFIIAegct+06DFYd2NVx6cUr83H6eWaMddS8RRLzZwIKkuNlxQZJQiIEKKFpakRlbS41I5VOTdmgamnWJZlT2ksVKeBqWBUVr8SMUpSbqRujhxTRA/5bOVnx5ceg97Vy3PrQsMytnJeoU3LEZeRjEUrCsGlsddCgRxnKbVNgnftQhE3ixbPPM1ClXbOtUXXRRxjnutUaVfVqpsdLDrI1y6okgbI6Xji2JPvyL2g06oiqHIUdX6ahY8sPBYRLKcuaWecCFVaVPOqne3S9dr9op8rUpvJq3upzEctTqpxaOfwScJGcXc1wqamfe7o79bzqhHc1Bh4iuhdrisU/I6Dq61FpIW2lZ8qR552MwkniaM3ohJ9YKKPYRiGYRiGYTxbPB3Hz0sAfAuA3yKiX8/bXgvgm4jojyGlev0+gL/5tE7YLM0OzMG8pN/o1BAJprXoE6JDdLObohwnB+d6tZ8iJmXHxZjf69ekAj5J/ZLjHo59CT59Po4IEnIO+b5e+Um7Tto0M+fmoKtaklkFsyIGSPsn50owpn81d+CdlCEJTqV9bU0OOVcbTB5PfQlqJWVGp+6MSuAq7g431y5ZcmdpwQ6Yg3vZogNpuQ7pl6DGdlQpN14JAofjUF2PHL9Kncpt0uMw5fELam7Jal1l7rFDCLUQ2LmIMdYuBjlfOx91f5aUmdwucWEF5U7T/SbBtvSrDvTbOS39FqIrdVxk3zE6BErimZ4/gV2pe9XOg3Kduc+9SvVphU/Zp70PxYkUoyt1q4oYmgUI0vsxVatvVS4v3q21JPevV/3QCp5CK0pol5xGxBJdR0u7wKLqaxEp2rpeJ4kx2h2m585S2mvb3o5iEcv0vNbPix0Rheb2ybi3S6i3go+8F5ejzKfkgHOz2EzLAs3SM2UJ3b8nOYkMwzAMwzAMw7h9PJ1Vvf4PFtfswTtu5cTFoVF+jZ8Fhq1y+3QuwlNEZLUCjnKLOKBaDl3SrCS4K8tvA5WwAMwOFgmallZSOtePpdbLlAvmSjDOSEF470OV0iPHLME2UH45B+oUl5DbKMcX8WHlp8otwUygfAxdZFULIuX68ntdx6YtRKwDzSJGqBSuNu3D5XGQMQJQ0tk8RfT9XB8HQCm+TVyvOJU+5B1hR8awiFpuTrVxmGtA6eDYU8RqmKoiw3INY/SlXpIg4xvUmGkRoMt9LteYttWiFQB4mgsi6xovvhEceh92lnD3vn4v49QKK7KXzF9Swpium1M5TTALdbqQuKQTSb/IGImbTfeNHnfpy45iqbmlx7GjiL6bxVppQ+R59b62ZpOMbyQqjrDN1FVC76qb5nRJ9R1Z+Q5Iq+2tu6mqC+VdrO5pYbXgoHNIri9x14ljqBSQz/W24OsV6EraYz6WbqeMjUafd4ou1f3Ky8NPcJVjUOa6uOnkuSRF3LVw3AombR0teZ4NLpTnRhFF2RVXXCuOdhQBPwtBfZ5/UfXp0vdOcnJWfaHGRp7tgtxL5voxDMMwDMMwjNvPTa3qdavo5dx1QA6g/GoPSIpLVEH+7mpZ+nvAXABYF4fVjpsQZ8eMnCeoY2iBJWR3iwQ0I/lyjDGmxKayP1B+IZd2tEGZroMSOYk1gQlexK2cijO4gCn/4r7N+2tHiNQzAlC5P1bdVKfZFPEgB6PiqmnSTGQ/CcB3VlxqnCgidozBFxfWZuqw6iYEBo7GfsdtIfV12rZJ6o5uk/Tpln0KZCk5FuRanTp2p1wem9DNLid2xQ0lqXDaMZPqqbjKvSUOoC7X9iHi6ryyn8ecBiYiQKecFG2fTdFh5HnutI4OeS/pUm2/b6au9I5X94x2KImzC0hCjBRtnqLDUegXx9F1szNO0tFEWJG+EjEmiZOuclmVa1PzXoterM+n7oviSqG8gt3UYZWFBrkHRLDyLoI4iYdp9bPZrZQOWAuT+vzSTyLCaBecdiNJ31SCEKmi1TmlUNqghTzHs/CpU+VKG/IYtEu0b0JXFeqWVNDZCVeLPq04ieb+kn7VzqPizGHCRn0/BFfXO8vt0svZEzE6xCrVLjJhlUVA5+oaVS36uacFIT23dRrs9YpGG4ZhGIZhGIZxezhV4QeYgzEJdnUQIcGlBBltcNEG1bIctnxXUm5kn66biuii3T4SXFar+DBVyzX7HIhJ0HSc02LW3VSEhd7PbZRVptrC1N5HjOQr8af3Ab06pwRbQK4v4pKooFOrpF6IBI0SFK67uW5RcSuo+h0JV5b1luOVYBSMvgtZ/EEVeAIpiGZijFkQS9caSxCtVybb68dyTKAu3qodHZJ+BqCqSyRLaIvwNbiAwYVSj6edC2P08KzSV4As2tVpYbIqm17qW8QuqVEkLgwZQ0fjjrsDQAn8e5fStWS8N6Gbj6XSkqQPZQ5WwTC4chEVdwwlh9jKTyXlTwulQaUoaseRtKNT16zPresZSb8ejENpl/S9Lkis7xPd7pWfqtQjQUSbTeh2Un90IWXvInwfMfgAn+9lOTaAUjupo4hVP1W1i/R9KQ7BsVl9S9f/0Y4bua9W3VRW0ZPjSUpazIJYZMKeH0u/AihCrAgrkQkbzPdU5SBzMldmx2CvnDkeXI0RUNeqkmknYyZpnTL3Na0bbt2Npb0+F/PufCzuPzmurumlx6anOeWr67S4vpzOVc7fOqu0EKRWQpSx00WrDcMwDMMwDMN4djhV4YfAVUpC+yu9BCJLQb5eEWntJxxNPaIKeuUXfTl+KfSq3BElMAKwZb/jAGrTH8bosM2CDxGX1BKpGy2/Xk8xrbokDiAdxogzRoJWfe29Cjb1NYubRYSWMadFhegwqF/sex/QIaZVwfSyzqgFMhGYtNgjyOpkOoivUjayW0UHtZLeomvAOOJSgBqYHQ/iymFSKWAqSNcrbkk6Ue9TitbBOICBOd0kj68e25VPYyIOKS0waZeJzCF5HZjgo0P0de0cAOhjI9DkY8pYyXwS8WV2M/i59hIx+hx86zktKWgyXyWg1ylCY/DFcSP9xUBxKfVqTmuhcxt9cV3J3PG5P8fg0XWxCBdj9NjIOARfpatpQakUSZf5IkIPz+eXlotoS3EW0KTQskbcd0ByNbUFr+VYkZMgN4UOm1A7okKYr733obp/2meACA56fsr16VpHgQlDFr32N+udOaBXLdNCsuwjNXG0YC2ClS5gPUZfBA9mLituadFLxCJJZSurfWVBTNos7Sl9m9/rOX0U+lKLzKmUvinMdYOG3DdT7DC4UAnquqhzOedSSpZ68LUCkFy7RtceWxKTDMMwDMMwDMO4PZy640c7UaZQr54kv2LHWLt9SnpADvS20Ver8OgixFJvQwc+OmVhVEGnTv8ov37nQKSIHN1UAm75BV9+QZdaHLquUJvSEzmtJtT7UFKyxLkj6S5btTKUDvIH5cRo04Pk9SbOQ1ilwuR2juxLYNwTV+ITAKzc7CDaG8ZKKJE+EoFN+lTSStqV0a5tV9W46eB4VCKQfOaJy4pGkv4XIMGxq44lTqcUvOaANDtRpBB1xOyMEUFCCz5y3t5FOOXGWXVTVajbEZf0nU3oynitPDC4CTGn+DgwjrPDRqeMeRUUi9Cyneb5OsWumnuS5ifXJ312rt9i8AHb4HdW29KCU1V7qc9/8/VrEUCcUyLordX9I2PFPBer1gG5iCZVOhIxghJlxQkmLjYpXs08iwC6+LQITACw142VI23lpuq6tXgSkRw5604JHNq9p5xv8n4KXXnOFIeZixgopW8Nqg6Qvn/FwSRzRdDnY6ZZRJNt0ub8V9IXZT7Is4OI0dPc995FeJyQ2tU4x47GPs8xKumnInDJuK3U9YqDTGrpVEXQQcWNtCTCSF/qmmktxa3E83u5Rl3rSb+WMUnXYWlfhmEYhmEYhvFscLqOn5zC0ooxOq1EFzoGsjjj5l/aU50dXwkOU3TY68eyDLwWcCSlzAOL6V0SBIn7QtBpCdJWcaroAs/SDl2oWKd96CBKAjK57qOpL8G0BMG6CDQwp7/1bg6SHRhdDtKm6Oulupmq69JLYuv0E+186ChgGztsg98Rd3SdGynKrPtEXgNzsWBdIHliNxfQxSxs6HaJA0UXeF35qRa4sgAjzghmKuPdUSxpLEAS4qT9OmWoTW2Z2BUBrsvB98G4KiLkyk3Y61LKl6QubairBILU9l0RRtjrxp0CytJGcbqI80yuVQpdf+rwPIZuwrk8t3WfdpjdNCXdbOqwl11G0qfSvz7Py1KjJqe3SXFyT3OdKwBY+2k3FVKN43Ho4H0s59PzPjb3gozzFF05T1ktSt37Os1QL0XuqRYjiLNTiutC1zL3tBCn2yHze5P7RcamuHJcANgBNKduahFUxkvTuYgui4Kyj+ApYu2nSthduYDe0c6YixNR36tDFmWnnLpYij/nuXe+31YCoBxLX6PPdZ9EzBzy3AKAFeZi5o523ZKDmzBFj2OuV0qLC+fU/RyUEF/mgOqv9nvlWRJrUcgwDMMwDMMwjNvDqQo/niIuDUe1I0Y7gHIqkwgbUdwS0WOKviz5Lt8Zgy81J1KwmF5LPZoxOvioVl7qphKAlsBPhJs41+4oS8hzrFbw8uWX6fxrPeZgt0P9y7YEZJLCBsypQiJMhUb4aovlyspXBOB8v62L1qqCxpyFEZ1yJekbks6iV1HTy2pLMWXZLuffw1iJG2XMHJW6KgBKUCniiCCOm95FOJWeJmIEByppXswpdSXkMS3OKheL00BEsTmVKrX1OPRl3syOkYghx+dyLEl5E0eCOC3EdTW5euW3TehwNPXY68biCHLE1TXK6lR9rkWk57IE7rMAFzBlp5q0axs8Oh8QnaucJohpvDsXcTx1OBz7MjdFUBR3TVpe3pf74cKwKXNOu8pk9TUZ46OpL+4bSfE5HF2VkiZ9LYLV4EJxvs3pSjNaIPMuVs4fESdlNa3IhJ5CNQfLXFDpj9qxFZmKa24MHtssHvVuFv0A4Hia70Fpg9wD8rzwNPeLiDzb2BVRRAo66wLgMq7aVdY6/PTcBRwOpqGqS6TT/to0J7mvzvfbnZQxR4y1H8tYFQcYowgt8jxo63sxp9UDN83j/mjsU7qoiynlTIl06Xk6i8oi1A0uwPm6CLm0b36jagKpa5R7SPpNF5rWLk3DMAzDMAzDMG4vpyr8TNHh8aMLRdSQILgNrErjKGKrUhMkCDsYB1zbrLCZPGIOzoYuIDLQ+4ijbV+7XHxEiElc8DlIJKCIDb2LaWWq6LAJHtvJlyCq9ynQJaQAazt1iAzs9VP5xX0Tuirwl9WYCCh/xXkkwpKjOZVLB/Xrbqzq+IhgsQld2UevJhbinJLlXXI4HaIv35c0NUBcJqpOC3K9FHbo3OywmaLDlc26KkitRS29utQ2+CJC7B+vKidD6vfUp14JTgGqyHauA5NSXFBWmlp1EzZTKj4sgaJOu5OxPRz7UodJhALtJJGV03SKSlRB+OE4VEWrdb93LuLadlW2Sa0ceS0im7i3dLsiE65u12UVO71dO2+K0EkBEUkAklXKVnl5+eOpq9KxCMAmzHNUixH7m3Vp/xgcop5vXUDXrOq0ViuKsQhJwePqdkCIhHtWtatkEzqMMfX7ZuzgHWPopqqaCwNAdLgy9qVdIabAvu9CGQ8J9kcRfNX8C0xw0eF47OBodoqJ2FOcaZPHhmdBVcZX5oMsf56uMaVu9i5iw50a99l9JIJp7+Zi2OJ+kvu070IRXs7120X3i7gCRSw5mnocb4cy96TPRUwuRZuDx35cY5S0LhfLM+rIzQ7BqK61fdboIuJAEu5CdCl1Lt9HwLzaoPwN21mELrKRErQ8xyIquuZ/S2g3V9qAkr7XOuOA3ZQwwzAMwzAMwzBuD7ck/BDRywD8KFIm1Y8z8w9cb38JBBwxXK470wo+EthM7ErtE/mVf3ATJvapoG/vECIhAEXMGVxMgdSQzqcL4QJz8V4tBIiTBADW3Yh1N2LTJZHleErOBEkRiQzEHPhN0WF/uypug03ocLAdcDzWThzv0qpEm6mrAuTAhOPtgCk4DF12mYwdPrG9B5SFlqEL6Lu6/se6S4JTBMEPmypNzBFjM3UlVQhIzhJJGdFOH/1r/eA2yUnAHi4X4CbMwToADF0owfdeP5a0LiClhozB4+J6UxXY1ktMt6ldfRZNkrtjTvE6mnrsb1fAdlUC4yLo0LzylISavYsYVpsSgOq6J55Seo9ugyPGuX5b+lSuQ9J4xuirtKdN6IoT63jq0rLqTIhTh1EHxS4tP56KaEdEOAz9VMQaXTOnpKw1wa+jCIcUHO8frcu46uXsex9xfthi5SdsQoc+143xFHGcxaeOpGh4h+gYK5qKIyrkMVh3E46nNGdnMSEJmtI3Mi/LnFEOsHU3lWLnyYUyO1/ERbTupjIfROQ7mnp4iqWosncRF/otVt1Uxmq1Ssfd60YcD11x6+g5dxy6SiBriy3rYvDyfUEcQNqJI2mg8hwKnFbgOpr64lLr87NI0gJj5/DE0fkyplogk7TVMXhEl85xXmp2gUuNMl2HqFPzToQgeT5J6p7ct7JKGzALO7oPZFwgTrl+N/2t7JddT7VbLZRUP3mu3DMcV0KPdgLqYuqVCkjz/EEERnnG5LbKuIhwrotWG4ZhGIZhGIZxe3jG/8omIg/gxwD8WQCPAfhlIno7M7/3pO9o14KuldHWi5E0EGBOmZHPHh6u4ij0OAo93AUu9VokNcxRRMyi0bw8d8C1cYUr272SQhHk128/rxIlRYPvGTboKODquC77jzmwESFJgkoRLACg3wuYVnMQJtdUgrBMCRSRlodf5WXn9/oR42pb+qa4Q1wsjgipEySFmnsVTF7oNzikHhP7qljtwbhKjqM4r2Yk7Vj7uUbLA8MBAhOm6PH8i5uSArMJdQFqAFW9F0lj0jVQJKg9Dh2eOLpQgnq5pnWXhIvjqSspc70P2OtGbDdrbCePVZ8cPveujrH2Y2n74TQUwevicAwA2KoULL2SlASlEjRLweEkmCRRT1KfdNHbjiLggHPdIQYfcN5vy1zTDgdHSWzchA4jO2xjlwPiWYyboseEXFRcxAiVtjf4UMSEzkU8dO4qLvTbIg6kulZ5f6AIBMBU1Vw612/LEu0yR/dyn+pV20TQmtMbCY6Ai+tNcrvkOVUEjnxfTTkdS9K+jkOHkXyZf0Au0ozUr/P1z4WpH1xfAwAcTgMeXF9D5wKOQxLrznXbMkeOQ4/BT3D5Wtd+LO4YIDlqjkOHw2lIaXmbWfAr9yinQsciNAkrl1bDO99FOIo4nIZSdFlEqicPz1fOHALgh7luViKCgWput46VVTdhz4XSL11ObzuchjLfHMWSZhZV/TMROQFUNcXk/UEYkitR0vfGvjjVOhfLM+hwHErKoq6t1GU30jb6UmNLhLAx+p3V5qbo4dS9JX3TOn+WHFDJ4Sg1i/LKYhTL8SJTqq9kGIZhGIZhGMZt51Z+Xv1KAB9g5g8CABG9BcDLAZwo/EQQrm1Xc+FmoAqCAbUse6490fuADZKL5Rxt8d4nH8bF1TGef/5TOJjmNJyVnzC4CR85uA8X+g1cdgLc26eaQhe7DToXsb9dl/okIhwMLuB8t8HF7hj3dYeIIHx8cy+eyCsquZzusefH8iv4Pd0xHl1/GpEdrkx7uOA35Rovd1dxzm3gifGhzYP48NEDRYgS8WDPj7jUH+Khfr84P867DRwiDuMKh3GFq2GNTexwGIcihG1ih03sShC857bY8yOOQo9rYYXDvKx8iK4Ujxa3Ru8c1khCjwTl2+AR2WHwE/7fwb3pmN2IwU+41B/BUcT9fcAm9hjZYYrJcXWpO8QFv0lOk9jjfdcexv3DYXJmuBF7PhUkvjLt4cH1NWxjV1w40g+hjEvAvd0R7usOENnhifECNrHD/rRXBY6RCVfGNR45dwV7fsTIDu978mFc3ruG+1aHuQD0HKyv/YSVm0qfRyTXhIxjYMJT23PYTB2ee24f57st9vwWE3s8OFzFtWmFTUyCwlHo0VHEnt9WQsj+uMb+uC4BvHZVbYPHNs5pSBO7JGj5Cef7TdlPB8qDC3jp/b+Dh/srOI49Pry9jE9u78FRHCqH1f64xuHUIzqHi8MRIjschw5rP2K9lwLsw2lABCXRjCL2fKppdWVcl/ZN7HFpOMTFboM9vy1zzBHDI2LlJjxn2C/tG9njynQOv3nlETxy7go6F8pc1EKRHEvq+AhHoUeAw8pN+NDBA/jgkw9gyCveSbFlqYtz6cIhLq2P0FPAyGmeirjWu4CL/XERfeU8cp/IuT1FrNwIn8fHg3HMHTaxx5PjOXzs+F4cTkkYWfuxCEyX19fw1PYcLg2H6CniodV+me+OIsY8rp9//nHsT3ulXcImdOhcwJ4fscpplCuXakV9YnNxFno4Fc2+NBwlcTqn6m0nj97P9al6HwAmHE8dzg9b/KHzT8FTKkIuz5ZVfkbtuS0e6A9SOzg5pmT1v5WbSlH0C36DkT3+7+FlfOzwXozBlzZJCmmMDp2PuLR3VK5N2gygCKGRkpAzwVUuppIqlrcdZ2fTOo91Ed3zCnmGYRiGYRiGYdx+buVf2o8A+AP1/jEAL253IqJXAXgVAPQX7sOVt38O2KUfd2MHcAfIj+TsABDgJuAg//UbwB8zuiOgP4qYVoRra8L7//AL0H/RPpyLCMHh8Kk9dE/0mC6PuOeBgyQYjT0OHj8HRALWAY8890k859zVWUTY7uH3PvUgrl7bQ9jvQROlfYlB923xnMv7WPmU4qTTiyTQ/bmnvgjbgwFgwA0Bq70RRIzD/TUwOYAYFx44xCP3XimpNlN0OJ56XN0M2Iw9jg4GxOMOUiWXVhHdakI/TFj1KVUmxiQYhOAQI4GVUDYddqBjD/YMCmllHDCBIgAGugPCsE/oDoD+kOG3DDcyKAIUgWmP0n7HjLBK37vGjCe/0OPo+Vt05yZ0fcB6GLHqU8rJ4WbAwZU1cOTRXfPwR4TNc0ecv3yI3gfsX90DfXwFEMDP2eDBB67i/LAtKzdd3azwqacuYLoyoL+Svs8OCHuM6Z6A9eUj3HPuGOf6sdS4OdgMODxcITw1oH/Kwx8T3AQcPRLw8fsvghxj++Qa6CJoSEux+y6g6yK6LokTqz6lJz15uIdr19aIhx1o48BDhDv0WH3agR0wXAGOLzO2nzOCHIMnArYONBH8kUvz8ojgt4Db5Hl6zOgPGd2GQZHBRIgdMO05RA/EAXneEygw/AYgTmPUH0a4gFKk5Tc+70U4ejgi3BPQX9xgbz2WVDupM3S07bHddhiPemC/AwVCvDiB+gjXRdx38RDnhy16H3B1u8LjVy/g4OPn4Q884BjrF1zFF1z+JAY34ROHF/H7U4/jqcOnr5wHmDCsRozbDtP+AH/Vw2+R+nyb7slrX3GE516+Urn1AhOeuHIBzEDXRWy3HcLGp3sqEhAIcIzhEx1WTxHGe4DNgwG0JayecFg9mebh5hJhe4nxkQvPRffIIWJwiIEQJwcce7iNQ1xFDPcf4/zeBt4xpuDK0vLj6BEjIQaPGCgXEgdADN56uAOPfp/gNgQ4YNpj+C2huwYMV9Oz5uqjhOkCY3xwBA0RPDq4K126b6b0PfYM/7xDDMOEoZuwl1P7xO12NPW4thmwf/UcYkirE3Ik8NYBkwMFAo2E/orDsJ/G/9rnRlBI9yY7gEKa5xSRVr3ywCfviTj4vCGluKq0KwawnTocbnocHazAkwP5dC/EyQFHqe+6a4T+WrrXKQDTOWA6x4grLs8/N6VniZvS8+RgC+w9nvrGbxkucHnGgAAKXBca0nCa9wBwfJ9DWAH9NZR7RTIeP/VFHo++9CNP778+hmEYhmEYhmE8bYh5uTDnDb9I9I0AXsbMfz2//xYAL2bmV1/nO1cBvP8ZndA4i1wG8MSdboTxGYPNB0Nj88EQbC4YGpsPhsbmg6Gx+WAId+tc+FxmfnDpg1tx/HwUwKPq/fPytuvxfmb+ils4p3GGIKJfsflgCDYfDI3NB0OwuWBobD4YGpsPhsbmgyHYXNjlVqpp/jKAFxLRC4hoAPAKAG+/Pc0yDMMwDMMwDMMwDMMwbpVn7Phh5omIXg3g55CWc38jM7/ntrXMMAzDMAzDMAzDMAzDuCVuaRkVZn4HgHfcxFdefyvnM84cNh8Mjc0HQ2PzwRBsLhgamw+GxuaDobH5YAg2FxqecXFnwzAMwzAMwzAMwzAM4zObW6nxYxiGYRiGYRiGYRiGYXwGY8KPYRiGYRiGYRiGYRjGGeXUhB8iehkRvZ+IPkBE33ta5zXuDET0KBH9AhG9l4jeQ0TflbffT0TvJKLfy3/vy9uJiP5Vnh+/SURffmevwHg2ICJPRL9GRD+b37+AiN6Vx/0/5BUCQUSr/P4D+fPn39GGG7cdIrpERG8lot8hovcR0Z+w58PdCxH93fzfit8mop8iorU9H+4eiOiNRPRJIvptte2mnwdE9Mq8/+8R0SvvxLUYt8YJc+Ff5P9W/CYR/UciuqQ+e02eC+8noq9T2y3uOAMszQf12fcQERPR5fzeng1nnJPmAxF9Z35GvIeI/rnabs8HxakIP0TkAfwYgD8H4EUAvomIXnQa5zbuGBOA72HmFwH4KgDfkcf8ewH8PDO/EMDP5/dAmhsvzP97FYDXnX6TjVPguwC8T73/ZwB+mJk/H8CTAL49b/92AE/m7T+c9zPOFj8K4L8x8x8B8KVI88KeD3chRPQIgL8D4CuY+UuQVgp9Bez5cDfxJgAva7bd1POAiO4H8H0AXgzgKwF8n4hFxmcVb8LuXHgngC9h5j8K4HcBvAYA8r8rXwHgi/N3/nX+gcnijrPDm7A7H0BEjwL4WgAfUZvt2XD2eROa+UBEXw3g5QC+lJm/GMAP5u32fGg4LcfPVwL4ADN/kJm3AN6CNEDGGYWZP8bM786vryIFdY8gjftP5N1+AsBfzK9fDuAnOfGLAC4R0XNPt9XGswkRPQ/AXwDw4/k9AXgpgLfmXdr5IPPkrQC+Ju9vnAGI6F4AfwrAGwCAmbfM/BTs+XA30wHYI6IOwDkAH4M9H+4amPl/A/h0s/lmnwdfB+CdzPxpZn4SSSzYCRiNz2yW5gIz/3dmnvLbXwTwvPz65QDewswbZv4QgA8gxRwWd5wRTng2AEn0//sA9CpF9mw445wwH/42gB9g5k3e55N5uz0fGk5L+HkEwB+o94/lbcZdQLbhfxmAdwF4iJk/lj/6OICH8mubI2efH0H6j3TM7x8A8JT6x5we8zIf8udX8v7G2eAFAB4H8O8opf79OBGdhz0f7kqY+aNIv9B9BEnwuQLgV2HPh7udm30e2HPi7uDbAPzX/Nrmwl0IEb0cwEeZ+Teaj2w+3J18AYA/mVO//xcR/fG83eZDgxV3Np5ViOgCgJ8B8N3MvK8/Y2ZGrdQbZxQi+noAn2TmX73TbTE+I+gAfDmA1zHzlwE4wJzGAcCeD3cT2XL/ciRB8HMAnIf9Gmso7HlgAAAR/UOkUgJvvtNtMe4MRHQOwGsB/KM73RbjM4YOwP1IpUX+HoCfNhfwMqcl/HwUwKPq/fPyNuMMQ0Q9kujzZmZ+W978CUnRyH/Fjmdz5GzzEgDfQES/j2SpfClSjZdLObUDqMe8zIf8+b0APnWaDTaeVR4D8Bgzvyu/fyuSEGTPh7uTPwPgQ8z8ODOPAN6G9Myw58Pdzc0+D+w5cYYhor8K4OsBfHMWAgGbC3cjn4f0I8Fv5H9TPg/Au4noYdh8uFt5DMDbcorfLyFlFlyGzYcdTkv4+WUAL6S0QseAVGjp7ad0buMOkJXWNwB4HzP/kPro7QCkmv4rAfxntf1bc0X+rwJwRVm8jc9ymPk1zPw8Zn4+0v3/P5n5mwH8AoBvzLu180HmyTfm/e3X3jMCM38cwB8Q0RfmTV8D4L2w58PdykcAfBURncv/7ZD5YM+Hu5ubfR78HICvJaL7sovsa/M247McInoZUqr4NzDzofro7QBeQWmlvxcgFfX9JVjccWZh5t9i5ucw8/PzvykfA/Dl+d8V9my4O/lPAL4aAIjoCwAMAJ6APR926G68y63DzBMRvRrpJvMA3sjM7zmNcxt3jJcA+BYAv0VEv563vRbADyBZ8L4dwIcB/OX82TsA/HmkwluHAP7aqbbWuFP8AwBvIaJ/CuDXkIv95r//nog+gFTE7RV3qH3Gs8d3Anhz/o/uB5HueQd7Ptx1MPO7iOitAN6NlMbxawBeD+C/wJ4PdwVE9FMA/jSAy0T0GNIKPDf17wVm/jQRfT/SP+oB4J8w81JRWOMzmBPmwmsArAC8M2dw/CIz/y1mfg8R/TSSUDwB+A5mDvk4FnecAZbmAzO/4YTd7dlwxjnh+fBGAG+ktMT7FsAr849B9nxoIPuRzDAMwzAMwzAMwzAM42xixZ0NwzAMwzAMwzAMwzDOKCb8GIZhGIZhGIZhGIZhnFFM+DEMwzAMwzAMwzAMwzijmPBjGIZhGIZhGIZhGIZxRjHhxzAMwzAMwzAMwzAM44xiwo9hGIZhGIZhGIZhGMYZxYQfwzAMwzAMwzAMwzCMM8r/B5JVALvAKAWUAAAAAElFTkSuQmCC\n",
      "text/plain": [
       "<Figure size 1440x360 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "plt.figure(figsize=(20,5))\n",
    "plt.imshow(b['x'][3], origin='lower')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3f4a1197",
   "metadata": {},
   "source": [
    "## Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 174,
   "id": "77a28f1f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "951581"
      ]
     },
     "execution_count": 174,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import src.modules\n",
    "reload(src.modules)\n",
    "from src.modules import FactorizedConv, ResBlock, StreamBlock\n",
    "\n",
    "import src.model\n",
    "reload(src.model)\n",
    "from src.model import ASRModel\n",
    "\n",
    "config = {\n",
    "    'in_channel' : 40,\n",
    "    'd_model' : 128,\n",
    "    'n_vocab' : tokenizer.vocab_size,\n",
    "    'dropout' : 0.1\n",
    "}\n",
    "\n",
    "model = ASRModel(config, tokenizer)\n",
    "model.count_parameters()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 171,
   "id": "c084daa4",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "ASRModel(\n",
       "  (criterion): CTCLoss()\n",
       "  (streams): ModuleList(\n",
       "    (0): StreamBlock(\n",
       "      (conv_blocks): ModuleList(\n",
       "        (0): ResBlock(\n",
       "          (conv): FactorizedConv(\n",
       "            (f1): Conv1d(40, 128, kernel_size=(2,), stride=(1,), padding=(1,))\n",
       "            (f2): Conv1d(128, 128, kernel_size=(2,), stride=(1,))\n",
       "          )\n",
       "          (post): Sequential(\n",
       "            (0): BatchNorm1d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (1): ReLU()\n",
       "            (2): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "          (scale): Conv1d(40, 128, kernel_size=(1,), stride=(1,))\n",
       "          (pool): AvgPool1d(kernel_size=(2,), stride=(2,), padding=(0,))\n",
       "        )\n",
       "        (1): ResBlock(\n",
       "          (conv): FactorizedConv(\n",
       "            (f1): Conv1d(128, 128, kernel_size=(2,), stride=(1,), padding=(1,))\n",
       "            (f2): Conv1d(128, 128, kernel_size=(2,), stride=(1,))\n",
       "          )\n",
       "          (post): Sequential(\n",
       "            (0): BatchNorm1d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (1): ReLU()\n",
       "            (2): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "          (scale): Conv1d(128, 128, kernel_size=(1,), stride=(1,))\n",
       "          (pool): AvgPool1d(kernel_size=(2,), stride=(2,), padding=(0,))\n",
       "        )\n",
       "        (2): ResBlock(\n",
       "          (conv): FactorizedConv(\n",
       "            (f1): Conv1d(128, 128, kernel_size=(2,), stride=(1,), padding=(1,))\n",
       "            (f2): Conv1d(128, 128, kernel_size=(2,), stride=(1,))\n",
       "          )\n",
       "          (post): Sequential(\n",
       "            (0): BatchNorm1d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (1): ReLU()\n",
       "            (2): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "          (scale): Conv1d(128, 128, kernel_size=(1,), stride=(1,))\n",
       "          (pool): AvgPool1d(kernel_size=(1,), stride=(1,), padding=(0,))\n",
       "        )\n",
       "      )\n",
       "      (attention): MultiheadAttention(\n",
       "        (out_proj): NonDynamicallyQuantizableLinear(in_features=128, out_features=128, bias=True)\n",
       "      )\n",
       "      (lnorm1): LayerNorm((128,), eps=1e-05, elementwise_affine=True)\n",
       "      (ff1): Linear(in_features=128, out_features=128, bias=True)\n",
       "      (ff2): Linear(in_features=128, out_features=128, bias=True)\n",
       "      (lnorm2): LayerNorm((128,), eps=1e-05, elementwise_affine=True)\n",
       "    )\n",
       "    (1): StreamBlock(\n",
       "      (conv_blocks): ModuleList(\n",
       "        (0): ResBlock(\n",
       "          (conv): FactorizedConv(\n",
       "            (f1): Conv1d(40, 128, kernel_size=(2,), stride=(1,), padding=(1,), dilation=(2,))\n",
       "            (f2): Conv1d(128, 128, kernel_size=(2,), stride=(1,), padding=(1,), dilation=(2,))\n",
       "          )\n",
       "          (post): Sequential(\n",
       "            (0): BatchNorm1d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (1): ReLU()\n",
       "            (2): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "          (scale): Conv1d(40, 128, kernel_size=(1,), stride=(1,))\n",
       "          (pool): AvgPool1d(kernel_size=(2,), stride=(2,), padding=(0,))\n",
       "        )\n",
       "        (1): ResBlock(\n",
       "          (conv): FactorizedConv(\n",
       "            (f1): Conv1d(128, 128, kernel_size=(2,), stride=(1,), padding=(1,), dilation=(2,))\n",
       "            (f2): Conv1d(128, 128, kernel_size=(2,), stride=(1,), padding=(1,), dilation=(2,))\n",
       "          )\n",
       "          (post): Sequential(\n",
       "            (0): BatchNorm1d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (1): ReLU()\n",
       "            (2): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "          (scale): Conv1d(128, 128, kernel_size=(1,), stride=(1,))\n",
       "          (pool): AvgPool1d(kernel_size=(2,), stride=(2,), padding=(0,))\n",
       "        )\n",
       "        (2): ResBlock(\n",
       "          (conv): FactorizedConv(\n",
       "            (f1): Conv1d(128, 128, kernel_size=(2,), stride=(1,), padding=(1,), dilation=(2,))\n",
       "            (f2): Conv1d(128, 128, kernel_size=(2,), stride=(1,), padding=(1,), dilation=(2,))\n",
       "          )\n",
       "          (post): Sequential(\n",
       "            (0): BatchNorm1d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (1): ReLU()\n",
       "            (2): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "          (scale): Conv1d(128, 128, kernel_size=(1,), stride=(1,))\n",
       "          (pool): AvgPool1d(kernel_size=(1,), stride=(1,), padding=(0,))\n",
       "        )\n",
       "      )\n",
       "      (attention): MultiheadAttention(\n",
       "        (out_proj): NonDynamicallyQuantizableLinear(in_features=128, out_features=128, bias=True)\n",
       "      )\n",
       "      (lnorm1): LayerNorm((128,), eps=1e-05, elementwise_affine=True)\n",
       "      (ff1): Linear(in_features=128, out_features=128, bias=True)\n",
       "      (ff2): Linear(in_features=128, out_features=128, bias=True)\n",
       "      (lnorm2): LayerNorm((128,), eps=1e-05, elementwise_affine=True)\n",
       "    )\n",
       "    (2): StreamBlock(\n",
       "      (conv_blocks): ModuleList(\n",
       "        (0): ResBlock(\n",
       "          (conv): FactorizedConv(\n",
       "            (f1): Conv1d(40, 128, kernel_size=(2,), stride=(1,), padding=(2,), dilation=(3,))\n",
       "            (f2): Conv1d(128, 128, kernel_size=(2,), stride=(1,), padding=(1,), dilation=(3,))\n",
       "          )\n",
       "          (post): Sequential(\n",
       "            (0): BatchNorm1d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (1): ReLU()\n",
       "            (2): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "          (scale): Conv1d(40, 128, kernel_size=(1,), stride=(1,))\n",
       "          (pool): AvgPool1d(kernel_size=(2,), stride=(2,), padding=(0,))\n",
       "        )\n",
       "        (1): ResBlock(\n",
       "          (conv): FactorizedConv(\n",
       "            (f1): Conv1d(128, 128, kernel_size=(2,), stride=(1,), padding=(2,), dilation=(3,))\n",
       "            (f2): Conv1d(128, 128, kernel_size=(2,), stride=(1,), padding=(1,), dilation=(3,))\n",
       "          )\n",
       "          (post): Sequential(\n",
       "            (0): BatchNorm1d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (1): ReLU()\n",
       "            (2): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "          (scale): Conv1d(128, 128, kernel_size=(1,), stride=(1,))\n",
       "          (pool): AvgPool1d(kernel_size=(2,), stride=(2,), padding=(0,))\n",
       "        )\n",
       "        (2): ResBlock(\n",
       "          (conv): FactorizedConv(\n",
       "            (f1): Conv1d(128, 128, kernel_size=(2,), stride=(1,), padding=(2,), dilation=(3,))\n",
       "            (f2): Conv1d(128, 128, kernel_size=(2,), stride=(1,), padding=(1,), dilation=(3,))\n",
       "          )\n",
       "          (post): Sequential(\n",
       "            (0): BatchNorm1d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (1): ReLU()\n",
       "            (2): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "          (scale): Conv1d(128, 128, kernel_size=(1,), stride=(1,))\n",
       "          (pool): AvgPool1d(kernel_size=(1,), stride=(1,), padding=(0,))\n",
       "        )\n",
       "      )\n",
       "      (attention): MultiheadAttention(\n",
       "        (out_proj): NonDynamicallyQuantizableLinear(in_features=128, out_features=128, bias=True)\n",
       "      )\n",
       "      (lnorm1): LayerNorm((128,), eps=1e-05, elementwise_affine=True)\n",
       "      (ff1): Linear(in_features=128, out_features=128, bias=True)\n",
       "      (ff2): Linear(in_features=128, out_features=128, bias=True)\n",
       "      (lnorm2): LayerNorm((128,), eps=1e-05, elementwise_affine=True)\n",
       "    )\n",
       "  )\n",
       "  (dense): Linear(in_features=384, out_features=29, bias=True)\n",
       ")"
      ]
     },
     "execution_count": 171,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "17d75361",
   "metadata": {},
   "source": [
    "## train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 175,
   "id": "54ba5be3",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "GPU available: False, used: False\n",
      "TPU available: False, using: 0 TPU cores\n",
      "IPU available: False, using: 0 IPUs\n"
     ]
    }
   ],
   "source": [
    "name = 'test'\n",
    "logger = TensorBoardLogger(save_dir='logs/', name=name)\n",
    "checkpoint = ModelCheckpoint(\n",
    "    dirpath=f'weights/{name}', \n",
    "    filename='{epoch}--{val_loss:.2f}', \n",
    "    monitor='val_loss',\n",
    "    mode='min',\n",
    "    save_top_k=5, \n",
    "    every_n_epochs = 1\n",
    ")\n",
    "\n",
    "trainer = Trainer(\n",
    "    benchmark=True, \n",
    "    gpus=0, \n",
    "    accumulate_grad_batches=1,\n",
    "    logger=logger, \n",
    "    max_epochs=1,\n",
    "    callbacks=[checkpoint]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 176,
   "id": "b1edc82a",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Missing logger folder: logs/test\n",
      "\n",
      "  | Name      | Type       | Params\n",
      "-----------------------------------------\n",
      "0 | criterion | CTCLoss    | 0     \n",
      "1 | streams   | ModuleList | 940 K \n",
      "2 | dense     | Linear     | 11.2 K\n",
      "-----------------------------------------\n",
      "951 K     Trainable params\n",
      "0         Non-trainable params\n",
      "951 K     Total params\n",
      "3.806     Total estimated model params size (MB)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation sanity check: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/soroush/.local/lib/python3.8/site-packages/pytorch_lightning/trainer/data_loading.py:132: UserWarning: The dataloader, val_dataloader 0, does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` (try 8 which is the number of cpus on this machine) in the `DataLoader` init to improve performance.\n",
      "  rank_zero_warn(\n",
      "Global seed set to 42\n",
      "/home/soroush/.local/lib/python3.8/site-packages/pytorch_lightning/trainer/data_loading.py:132: UserWarning: The dataloader, train_dataloader, does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` (try 8 which is the number of cpus on this machine) in the `DataLoader` init to improve performance.\n",
      "  rank_zero_warn(\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "947e3b7cd047494f89343eb4fd6b1fbb",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Training: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "ename": "AssertionError",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mAssertionError\u001b[0m                            Traceback (most recent call last)",
      "\u001b[0;32m/tmp/ipykernel_832/3881363385.py\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0mtrainer\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfit\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmodel\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtl\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mvl\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[0;32m~/.local/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py\u001b[0m in \u001b[0;36mfit\u001b[0;34m(self, model, train_dataloaders, val_dataloaders, datamodule, train_dataloader, ckpt_path)\u001b[0m\n\u001b[1;32m    738\u001b[0m             )\n\u001b[1;32m    739\u001b[0m             \u001b[0mtrain_dataloaders\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtrain_dataloader\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 740\u001b[0;31m         self._call_and_handle_interrupt(\n\u001b[0m\u001b[1;32m    741\u001b[0m             \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_fit_impl\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmodel\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtrain_dataloaders\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mval_dataloaders\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdatamodule\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mckpt_path\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    742\u001b[0m         )\n",
      "\u001b[0;32m~/.local/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py\u001b[0m in \u001b[0;36m_call_and_handle_interrupt\u001b[0;34m(self, trainer_fn, *args, **kwargs)\u001b[0m\n\u001b[1;32m    683\u001b[0m         \"\"\"\n\u001b[1;32m    684\u001b[0m         \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 685\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mtrainer_fn\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    686\u001b[0m         \u001b[0;31m# TODO: treat KeyboardInterrupt as BaseException (delete the code below) in v1.7\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    687\u001b[0m         \u001b[0;32mexcept\u001b[0m \u001b[0mKeyboardInterrupt\u001b[0m \u001b[0;32mas\u001b[0m \u001b[0mexception\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.local/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py\u001b[0m in \u001b[0;36m_fit_impl\u001b[0;34m(self, model, train_dataloaders, val_dataloaders, datamodule, ckpt_path)\u001b[0m\n\u001b[1;32m    775\u001b[0m         \u001b[0;31m# TODO: ckpt_path only in v1.7\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    776\u001b[0m         \u001b[0mckpt_path\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mckpt_path\u001b[0m \u001b[0;32mor\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mresume_from_checkpoint\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 777\u001b[0;31m         \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_run\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmodel\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mckpt_path\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mckpt_path\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    778\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    779\u001b[0m         \u001b[0;32massert\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstate\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstopped\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.local/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py\u001b[0m in \u001b[0;36m_run\u001b[0;34m(self, model, ckpt_path)\u001b[0m\n\u001b[1;32m   1197\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1198\u001b[0m         \u001b[0;31m# dispatch `start_training` or `start_evaluating` or `start_predicting`\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1199\u001b[0;31m         \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_dispatch\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1200\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1201\u001b[0m         \u001b[0;31m# plugin will finalized fitting (e.g. ddp_spawn will load trained model)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.local/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py\u001b[0m in \u001b[0;36m_dispatch\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m   1277\u001b[0m             \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtraining_type_plugin\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstart_predicting\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1278\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1279\u001b[0;31m             \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtraining_type_plugin\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstart_training\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1280\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1281\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mrun_stage\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.local/lib/python3.8/site-packages/pytorch_lightning/plugins/training_type/training_type_plugin.py\u001b[0m in \u001b[0;36mstart_training\u001b[0;34m(self, trainer)\u001b[0m\n\u001b[1;32m    200\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mstart_training\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtrainer\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0;34m\"pl.Trainer\"\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m->\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    201\u001b[0m         \u001b[0;31m# double dispatch to initiate the training loop\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 202\u001b[0;31m         \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_results\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtrainer\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrun_stage\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    203\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    204\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mstart_evaluating\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtrainer\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0;34m\"pl.Trainer\"\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m->\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.local/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py\u001b[0m in \u001b[0;36mrun_stage\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m   1287\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mpredicting\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1288\u001b[0m             \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_run_predict\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1289\u001b[0;31m         \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_run_train\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1290\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1291\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0m_pre_training_routine\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.local/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py\u001b[0m in \u001b[0;36m_run_train\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m   1317\u001b[0m         \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfit_loop\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtrainer\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1318\u001b[0m         \u001b[0;32mwith\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mautograd\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mset_detect_anomaly\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_detect_anomaly\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1319\u001b[0;31m             \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfit_loop\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrun\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1320\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1321\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0m_run_evaluate\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m->\u001b[0m \u001b[0m_EVALUATE_OUTPUT\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.local/lib/python3.8/site-packages/pytorch_lightning/loops/base.py\u001b[0m in \u001b[0;36mrun\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m    143\u001b[0m             \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    144\u001b[0m                 \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mon_advance_start\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 145\u001b[0;31m                 \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0madvance\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    146\u001b[0m                 \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mon_advance_end\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    147\u001b[0m                 \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrestarting\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;32mFalse\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.local/lib/python3.8/site-packages/pytorch_lightning/loops/fit_loop.py\u001b[0m in \u001b[0;36madvance\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    232\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    233\u001b[0m         \u001b[0;32mwith\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtrainer\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mprofiler\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mprofile\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\"run_training_epoch\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 234\u001b[0;31m             \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mepoch_loop\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrun\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdata_fetcher\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    235\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    236\u001b[0m             \u001b[0;31m# the global step is manually decreased here due to backwards compatibility with existing loggers\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.local/lib/python3.8/site-packages/pytorch_lightning/loops/base.py\u001b[0m in \u001b[0;36mrun\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m    143\u001b[0m             \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    144\u001b[0m                 \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mon_advance_start\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 145\u001b[0;31m                 \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0madvance\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    146\u001b[0m                 \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mon_advance_end\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    147\u001b[0m                 \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrestarting\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;32mFalse\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.local/lib/python3.8/site-packages/pytorch_lightning/loops/epoch/training_epoch_loop.py\u001b[0m in \u001b[0;36madvance\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m    191\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    192\u001b[0m             \u001b[0;32mwith\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtrainer\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mprofiler\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mprofile\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\"run_training_batch\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 193\u001b[0;31m                 \u001b[0mbatch_output\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbatch_loop\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrun\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mbatch\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mbatch_idx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    194\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    195\u001b[0m         \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbatch_progress\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mincrement_processed\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.local/lib/python3.8/site-packages/pytorch_lightning/loops/base.py\u001b[0m in \u001b[0;36mrun\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m    143\u001b[0m             \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    144\u001b[0m                 \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mon_advance_start\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 145\u001b[0;31m                 \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0madvance\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    146\u001b[0m                 \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mon_advance_end\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    147\u001b[0m                 \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrestarting\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;32mFalse\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.local/lib/python3.8/site-packages/pytorch_lightning/loops/batch/training_batch_loop.py\u001b[0m in \u001b[0;36madvance\u001b[0;34m(self, batch, batch_idx)\u001b[0m\n\u001b[1;32m     86\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtrainer\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mlightning_module\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mautomatic_optimization\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     87\u001b[0m             \u001b[0moptimizers\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0m_get_active_optimizers\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtrainer\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0moptimizers\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtrainer\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0moptimizer_frequencies\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mbatch_idx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 88\u001b[0;31m             \u001b[0moutputs\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0moptimizer_loop\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrun\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0msplit_batch\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0moptimizers\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mbatch_idx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     89\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     90\u001b[0m             \u001b[0moutputs\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmanual_loop\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrun\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0msplit_batch\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mbatch_idx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.local/lib/python3.8/site-packages/pytorch_lightning/loops/base.py\u001b[0m in \u001b[0;36mrun\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m    143\u001b[0m             \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    144\u001b[0m                 \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mon_advance_start\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 145\u001b[0;31m                 \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0madvance\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    146\u001b[0m                 \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mon_advance_end\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    147\u001b[0m                 \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrestarting\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;32mFalse\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.local/lib/python3.8/site-packages/pytorch_lightning/loops/optimization/optimizer_loop.py\u001b[0m in \u001b[0;36madvance\u001b[0;34m(self, batch, *args, **kwargs)\u001b[0m\n\u001b[1;32m    213\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    214\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0madvance\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mbatch\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mAny\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mAny\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mAny\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m->\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m:\u001b[0m  \u001b[0;31m# type: ignore[override]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 215\u001b[0;31m         result = self._run_optimization(\n\u001b[0m\u001b[1;32m    216\u001b[0m             \u001b[0mbatch\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    217\u001b[0m             \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_batch_idx\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.local/lib/python3.8/site-packages/pytorch_lightning/loops/optimization/optimizer_loop.py\u001b[0m in \u001b[0;36m_run_optimization\u001b[0;34m(self, split_batch, batch_idx, optimizer, opt_idx)\u001b[0m\n\u001b[1;32m    264\u001b[0m         \u001b[0;31m# gradient update with accumulated gradients\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    265\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 266\u001b[0;31m             \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_optimizer_step\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0moptimizer\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mopt_idx\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mbatch_idx\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mclosure\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    267\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    268\u001b[0m         \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mclosure\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mconsume_result\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.local/lib/python3.8/site-packages/pytorch_lightning/loops/optimization/optimizer_loop.py\u001b[0m in \u001b[0;36m_optimizer_step\u001b[0;34m(self, optimizer, opt_idx, batch_idx, train_step_and_backward_closure)\u001b[0m\n\u001b[1;32m    376\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    377\u001b[0m         \u001b[0;31m# model hook\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 378\u001b[0;31m         lightning_module.optimizer_step(\n\u001b[0m\u001b[1;32m    379\u001b[0m             \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtrainer\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcurrent_epoch\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    380\u001b[0m             \u001b[0mbatch_idx\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.local/lib/python3.8/site-packages/pytorch_lightning/core/lightning.py\u001b[0m in \u001b[0;36moptimizer_step\u001b[0;34m(self, epoch, batch_idx, optimizer, optimizer_idx, optimizer_closure, on_tpu, using_native_amp, using_lbfgs)\u001b[0m\n\u001b[1;32m   1650\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1651\u001b[0m         \"\"\"\n\u001b[0;32m-> 1652\u001b[0;31m         \u001b[0moptimizer\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstep\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mclosure\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0moptimizer_closure\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1653\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1654\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0moptimizer_zero_grad\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mepoch\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mint\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mbatch_idx\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mint\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0moptimizer\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mOptimizer\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0moptimizer_idx\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mint\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.local/lib/python3.8/site-packages/pytorch_lightning/core/optimizer.py\u001b[0m in \u001b[0;36mstep\u001b[0;34m(self, closure, **kwargs)\u001b[0m\n\u001b[1;32m    162\u001b[0m         \u001b[0;32massert\u001b[0m \u001b[0mtrainer\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    163\u001b[0m         \u001b[0;32mwith\u001b[0m \u001b[0mtrainer\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mprofiler\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mprofile\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mprofiler_action\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 164\u001b[0;31m             \u001b[0mtrainer\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0maccelerator\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0moptimizer_step\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_optimizer\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_optimizer_idx\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mclosure\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[0;32m~/.local/lib/python3.8/site-packages/pytorch_lightning/accelerators/accelerator.py\u001b[0m in \u001b[0;36moptimizer_step\u001b[0;34m(self, optimizer, opt_idx, closure, model, **kwargs)\u001b[0m\n\u001b[1;32m    337\u001b[0m         \"\"\"\n\u001b[1;32m    338\u001b[0m         \u001b[0mmodel\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mmodel\u001b[0m \u001b[0;32mor\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mlightning_module\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 339\u001b[0;31m         \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mprecision_plugin\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0moptimizer_step\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmodel\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0moptimizer\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mopt_idx\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mclosure\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    340\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    341\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0moptimizer_zero_grad\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcurrent_epoch\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mint\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mbatch_idx\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mint\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0moptimizer\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mOptimizer\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mopt_idx\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mint\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m->\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.local/lib/python3.8/site-packages/pytorch_lightning/plugins/precision/precision_plugin.py\u001b[0m in \u001b[0;36moptimizer_step\u001b[0;34m(self, model, optimizer, optimizer_idx, closure, **kwargs)\u001b[0m\n\u001b[1;32m    161\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0misinstance\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmodel\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mpl\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mLightningModule\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    162\u001b[0m             \u001b[0mclosure\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mpartial\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_wrap_closure\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmodel\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0moptimizer\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0moptimizer_idx\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mclosure\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 163\u001b[0;31m         \u001b[0moptimizer\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstep\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mclosure\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mclosure\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    164\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    165\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0m_track_grad_norm\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtrainer\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0;34m\"pl.Trainer\"\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m->\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.local/lib/python3.8/site-packages/torch/optim/optimizer.py\u001b[0m in \u001b[0;36mwrapper\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m     86\u001b[0m                 \u001b[0mprofile_name\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m\"Optimizer.step#{}.step\"\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mformat\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mobj\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m__class__\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m__name__\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     87\u001b[0m                 \u001b[0;32mwith\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mautograd\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mprofiler\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrecord_function\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mprofile_name\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 88\u001b[0;31m                     \u001b[0;32mreturn\u001b[0m \u001b[0mfunc\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     89\u001b[0m             \u001b[0;32mreturn\u001b[0m \u001b[0mwrapper\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     90\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.local/lib/python3.8/site-packages/torch/autograd/grad_mode.py\u001b[0m in \u001b[0;36mdecorate_context\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m     26\u001b[0m         \u001b[0;32mdef\u001b[0m \u001b[0mdecorate_context\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     27\u001b[0m             \u001b[0;32mwith\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m__class__\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 28\u001b[0;31m                 \u001b[0;32mreturn\u001b[0m \u001b[0mfunc\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     29\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0mcast\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mF\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdecorate_context\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     30\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.local/lib/python3.8/site-packages/torch/optim/adamw.py\u001b[0m in \u001b[0;36mstep\u001b[0;34m(self, closure)\u001b[0m\n\u001b[1;32m     90\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mclosure\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     91\u001b[0m             \u001b[0;32mwith\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0menable_grad\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 92\u001b[0;31m                 \u001b[0mloss\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mclosure\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     93\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     94\u001b[0m         \u001b[0;32mfor\u001b[0m \u001b[0mgroup\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mparam_groups\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.local/lib/python3.8/site-packages/pytorch_lightning/plugins/precision/precision_plugin.py\u001b[0m in \u001b[0;36m_wrap_closure\u001b[0;34m(self, model, optimizer, optimizer_idx, closure)\u001b[0m\n\u001b[1;32m    146\u001b[0m         \u001b[0mconsistent\u001b[0m \u001b[0;32mwith\u001b[0m \u001b[0mthe\u001b[0m\u001b[0;31m \u001b[0m\u001b[0;31m`\u001b[0m\u001b[0;31m`\u001b[0m\u001b[0mPrecisionPlugin\u001b[0m\u001b[0;31m`\u001b[0m\u001b[0;31m`\u001b[0m \u001b[0msubclasses\u001b[0m \u001b[0mthat\u001b[0m \u001b[0mcannot\u001b[0m \u001b[0;32mpass\u001b[0m\u001b[0;31m \u001b[0m\u001b[0;31m`\u001b[0m\u001b[0;31m`\u001b[0m\u001b[0moptimizer\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstep\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mclosure\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;31m`\u001b[0m\u001b[0;31m`\u001b[0m \u001b[0mdirectly\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    147\u001b[0m         \"\"\"\n\u001b[0;32m--> 148\u001b[0;31m         \u001b[0mclosure_result\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mclosure\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    149\u001b[0m         \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_after_closure\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmodel\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0moptimizer\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0moptimizer_idx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    150\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0mclosure_result\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.local/lib/python3.8/site-packages/pytorch_lightning/loops/optimization/optimizer_loop.py\u001b[0m in \u001b[0;36m__call__\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m    158\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    159\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0m__call__\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mAny\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mAny\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m->\u001b[0m \u001b[0mOptional\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mTensor\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 160\u001b[0;31m         \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_result\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mclosure\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    161\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_result\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mloss\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    162\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.local/lib/python3.8/site-packages/pytorch_lightning/loops/optimization/optimizer_loop.py\u001b[0m in \u001b[0;36mclosure\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m    140\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mclosure\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mAny\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mAny\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m->\u001b[0m \u001b[0mClosureResult\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    141\u001b[0m         \u001b[0;32mwith\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_profiler\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mprofile\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\"training_step_and_backward\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 142\u001b[0;31m             \u001b[0mstep_output\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_step_fn\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    143\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    144\u001b[0m             \u001b[0;32mif\u001b[0m \u001b[0mstep_output\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mclosure_loss\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.local/lib/python3.8/site-packages/pytorch_lightning/loops/optimization/optimizer_loop.py\u001b[0m in \u001b[0;36m_training_step\u001b[0;34m(self, split_batch, batch_idx, opt_idx)\u001b[0m\n\u001b[1;32m    433\u001b[0m             \u001b[0mlightning_module\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_current_fx_name\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m\"training_step\"\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    434\u001b[0m             \u001b[0;32mwith\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtrainer\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mprofiler\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mprofile\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\"training_step\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 435\u001b[0;31m                 \u001b[0mtraining_step_output\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtrainer\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0maccelerator\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtraining_step\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mstep_kwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    436\u001b[0m                 \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtrainer\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtraining_type_plugin\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mpost_training_step\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    437\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.local/lib/python3.8/site-packages/pytorch_lightning/accelerators/accelerator.py\u001b[0m in \u001b[0;36mtraining_step\u001b[0;34m(self, step_kwargs)\u001b[0m\n\u001b[1;32m    217\u001b[0m         \"\"\"\n\u001b[1;32m    218\u001b[0m         \u001b[0;32mwith\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mprecision_plugin\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtrain_step_context\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 219\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtraining_type_plugin\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtraining_step\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0mstep_kwargs\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mvalues\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    220\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    221\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mpost_training_step\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m->\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.local/lib/python3.8/site-packages/pytorch_lightning/plugins/training_type/training_type_plugin.py\u001b[0m in \u001b[0;36mtraining_step\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m    211\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    212\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mtraining_step\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 213\u001b[0;31m         \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmodel\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtraining_step\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    214\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    215\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mpost_training_step\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/mnt/d/Projects/ConvolutionalSpeechRecognition/src/model.py\u001b[0m in \u001b[0;36mtraining_step\u001b[0;34m(self, batch, batch_idx)\u001b[0m\n\u001b[1;32m     72\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     73\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mtraining_step\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mbatch\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mbatch_idx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 74\u001b[0;31m         \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstep\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mbatch\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     75\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     76\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mvalidation_step\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mbatch\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mbatch_idx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/mnt/d/Projects/ConvolutionalSpeechRecognition/src/model.py\u001b[0m in \u001b[0;36mstep\u001b[0;34m(self, batch, mode)\u001b[0m\n\u001b[1;32m     47\u001b[0m         \u001b[0;31m# log_probs shape (batch_size, output_len, num_classes)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     48\u001b[0m         \u001b[0mlog_probs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0morth_loss\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mforward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mbatch\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m'x'\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 49\u001b[0;31m         \u001b[0;32massert\u001b[0m \u001b[0mlog_probs\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mshape\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m>=\u001b[0m \u001b[0mbatch\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m'y'\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mshape\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     50\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     51\u001b[0m         \u001b[0;31m# CTC_Loss expects input shape\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mAssertionError\u001b[0m: "
     ]
    }
   ],
   "source": [
    "trainer.fit(model, tl, vl)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 177,
   "id": "faebae28",
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer.save_checkpoint(f'weights/{name}/last.ckpt')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "15b894ce",
   "metadata": {},
   "source": [
    "## predict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 182,
   "id": "c6c10284",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "h\n"
     ]
    }
   ],
   "source": [
    "# import librosa\n",
    "\n",
    "# x, _ = librosa.load('test.wav', sr=16000)\n",
    "# x = librosa.features.mfcc(y=x, n_mfcc=40)\n",
    "# x = torch.tensor(x)\n",
    "x = val_data[0][0]\n",
    "x = mfcc(x[0])\n",
    "print(model.transcribe(x))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "90a11398",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}