diff --git a/Quick Draw CNN.ipynb b/Quick Draw CNN.ipynb
index ae095a7..926650b 100644
--- a/Quick Draw CNN.ipynb	
+++ b/Quick Draw CNN.ipynb	
@@ -96,7 +96,7 @@
    ],
    "source": [
     "dataset_dir = \"quick_draw_dataset\"\n",
-    "files = os.listdir(dataset_dir)\n",
+    "files = [name for name in os.listdir(dataset_dir) if \".npy\" in name]\n",
     "max_size_per_cl = 1500\n",
     "draw_class = []\n",
     "\n",