Filter for files in the repository

thibo73800 · Mar 27, 2019 · f61e9f1 · f61e9f1
1 parent 06be2a2
commit f61e9f1
Showing 1 changed file with 1 addition and 1 deletion.
diff --git 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",