Fix tutorials test (#3458)

Summary: Pull Request resolved: #3458

Reviewed By: Cesar-Cardoso

Differential Revision: D70631834

fbshipit-source-id: 4e59b4df955103f0eca73907cb80d48c3965c0b3

Author: mpolson64

tutorials/automl/automl.ipynb

@@ -504,10 +504,10 @@
         "        )\n",
         "\n",
         "        batch_size = assert_is_instance(parameters[\"batch_size\"], int)\n",
+        "        num_epochs = len(train_x) // batch_size\n",
         "\n",
         "        start_time = time.time()\n",
-        "        is_early_stopped = False  # Flag to indicate if the trial was early stopped\n",
-        "        for i in range(0, len(train_x) // batch_size):\n",
+        "        for i in range(0, num_epochs):\n",
         "            start_idx = i * batch_size\n",
         "            end_idx = (i + 1) * batch_size\n",
         "\n",
@@ -520,26 +520,26 @@
         "                \"score\": clf.score(valid_x, valid_y),\n",
         "                \"training_time\": time.time() - start_time,\n",
         "            }\n",
+        "\n",
+        "            # On the final epoch call complete_trial and break, else call attach_data\n",
+        "            if i == num_epochs - 1:\n",
+        "                client.complete_trial(\n",
+        "                    trial_index=trial_index,\n",
+        "                    raw_data=raw_data,\n",
+        "                    progression=end_idx,  # Use the index of the last example in the batch as the progression value\n",
+        "                )\n",
+        "                break\n",
+        "\n",
         "            client.attach_data(\n",
         "                trial_index=trial_index,\n",
         "                raw_data=raw_data,\n",
-        "                progression=end_idx,  # Use the index of the last example in the batch as the progression value\n",
+        "                progression=end_idx,\n",
         "            )\n",
         "\n",
         "            # If the trial is underperforming, stop it\n",
         "            if client.should_stop_trial_early(trial_index=trial_index):\n",
-        "                client.mark_trial_early_stopped(\n",
-        "                    trial_index=trial_index,\n",
-        "                    raw_data=raw_data,\n",
-        "                    progression=end_idx,\n",
-        "                )\n",
-        "                is_early_stopped = True\n",
-        "                break\n",
-        "\n",
-        "        if not is_early_stopped:\n",
-        "            # If the trial was not early stopped, mark it as completed\n",
-        "            client.complete_trial(trial_index=trial_index)\n",
-        "            print(f\"Completed trial {trial_index} with {parameters=}, {raw_data=}\")"
+        "                client.mark_trial_early_stopped(trial_index=trial_index)\n",
+        "                break\n"
       ]
     },
     {

tutorials/early_stopping/early_stopping.ipynb

@@ -352,9 +352,7 @@
         "\n",
         "            # If the trial is underperforming, stop it\n",
         "            if client.should_stop_trial_early(trial_index=trial_index):\n",
-        "                client.mark_trial_early_stopped(\n",
-        "                    trial_index=trial_index, raw_data=raw_data, progression=t\n",
-        "                )\n",
+        "                client.mark_trial_early_stopped(trial_index=trial_index)\n",
         "                break"
       ]
     },