Refer to solution correctly in dqd tutorial vis

icaros-usc · Jan 3, 2024 · 771a72b · 771a72b
1 parent 4ce3d47
commit 771a72b
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/tutorials/tom_cruise_dqd.ipynb b/tutorials/tom_cruise_dqd.ipynb
@@ -950,7 +950,10 @@
     "    if not df_box.empty:\n",
     "        # Select the solution with highest objective in the box.\n",
     "        max_obj_idx = df_box['objective'].argmax()\n",
-    "        sol = df_box.loc[:, \"solution_0\":].iloc[max_obj_idx]\n",
+    "        # StyleGAN solutions have 7,168 dimensions, so the final solution col\n",
+    "        # is solution_7167.\n",
+    "        sol = df_box.loc[:, \"solution_0\":\"solution_7167\"].iloc[max_obj_idx]\n",
+    "\n",
     "\n",
     "        # Convert the latent vector solution to an image.\n",
     "        img = classifier.generate_image(np.array(sol))[0].cpu().detach()\n",