diff --git a/examples/frameworks/pytorch/notebooks/audio/audio_classifier_UrbanSound8K.ipynb b/examples/frameworks/pytorch/notebooks/audio/audio_classifier_UrbanSound8K.ipynb
index c1a67972..c22765b6 100644
--- a/examples/frameworks/pytorch/notebooks/audio/audio_classifier_UrbanSound8K.ipynb
+++ b/examples/frameworks/pytorch/notebooks/audio/audio_classifier_UrbanSound8K.ipynb
@@ -293,8 +293,8 @@
    "source": [
     "def test(model, epoch):\n",
     "    model.eval()\n",
-    "    class_correct = list(0. for i in range(10))\n",
-    "    class_total = list(0. for i in range(10))\n",
+    "    class_correct = list(0. for i in range(len(classes)))\n",
+    "    class_total = list(0. for i in range(len(classes)))\n",
     "    with torch.no_grad():\n",
     "        for idx, (sounds, sample_rate, inputs, labels) in enumerate(test_loader):\n",
     "            inputs = inputs.to(device)\n",