update demo

corajr · Feb 11, 2022 · ad5eec3 · ad5eec3
1 parent a3a2577
commit ad5eec3
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/demo.ipynb b/demo.ipynb
@@ -106,7 +106,10 @@
     "model = model.to(device)\n",
     "\n",
     "with torch.no_grad():\n",
-    "    caption = model.generate(image, sample=False, num_beams=3, max_length=20, min_length=5)\n",
+    "    # beam search\n",
+    "    caption = model.generate(image, sample=False, num_beams=3, max_length=20, min_length=5) \n",
+    "    # nucleus sampling\n",
+    "    # caption = model.generate(image, sample=True, top_p=0.9, max_length=20, min_length=5) \n",
     "    print('caption: '+caption[0])"
    ]
   },