Add assertion in demo

notebooks/Demo.ipynb

@@ -3,7 +3,7 @@
   {
    "cell_type": "code",
    "execution_count": 19,
-   "id": "9ef0e433",
+   "id": "27deb847",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -437,17 +437,17 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 42,
+   "execution_count": 46,
    "id": "32621820",
    "metadata": {},
    "outputs": [
     {
      "data": {
       "text/plain": [
-       "tensor([[1., 0., 0.]], grad_fn=<SoftmaxBackward0>)"
+       "0"
       ]
      },
-     "execution_count": 42,
+     "execution_count": 46,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -455,79 +455,30 @@
    "source": [
     "input = wav.unsqueeze(0)\n",
     "output = model(input)\n",
-    "prediction = train_dataset.labels[output]"
+    "prediction = torch.argmax(output, 1).item()\n",
+    "prediction"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 28,
+   "execution_count": 49,
    "id": "9455a818",
    "metadata": {},
    "outputs": [
     {
      "data": {
       "text/plain": [
-       "torch.Size([1, 128, 9, 31])"
-      ]
-     },
-     "execution_count": 28,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "prediction = torch.argmax(output, 1)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 34,
-   "id": "e23f2494",
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "torch.Size([1, 60032])"
+       "True"
       ]
      },
-     "execution_count": 34,
+     "execution_count": 49,
      "metadata": {},
      "output_type": "execute_result"
     }
    ],
    "source": [
-    "torch.nn.Flatten()(wav).shape"
+    "actual_output == prediction"
    ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 43,
-   "id": "246fd00b",
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "0"
-      ]
-     },
-     "execution_count": 43,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "actual_output"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "0dabc10c",
-   "metadata": {},
-   "outputs": [],
-   "source": []
   }
  ],
  "metadata": {