From dae1a1d11202f32b5a835c33ee633cf0d5fd1c8f Mon Sep 17 00:00:00 2001 From: Kumar Ashutosh Date: Fri, 22 Oct 2021 02:27:06 -0500 Subject: [PATCH] Minor edit when training with a custom dataset (#484) When we want to train a custom dataset, the code should have the flexibility to accommodate variable number of classes. --- .../notebooks/audio/audio_classifier_UrbanSound8K.ipynb | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) 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",